chalkcompute 2.1.0__tar.gz → 2.1.2__tar.gz

{chalkcompute-2.1.0 → chalkcompute-2.1.2}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: chalkcompute
-Version: 2.1.0
+Version: 2.1.2
 Requires-Dist: grpcio>=1.60.0
 Requires-Dist: protobuf>=4.25.0
 Requires-Dist: httpx>=0.25.0
 Requires-Dist: blake3>=1.0.0
-Requires-Dist: chalk-remote-call-python>=1.4.0
+Requires-Dist: chalk-remote-call-python>=1.6.3
 Requires-Dist: pyyaml>=6.0.0
 Requires-Dist: opentelemetry-api>=1.20.0
 Requires-Dist: opentelemetry-sdk>=1.20.0

{chalkcompute-2.1.0 → chalkcompute-2.1.2}/chalkcompute/_cls.py

@@ -47,6 +47,7 @@ import ast
 import inspect
 import logging
 import os
+import textwrap
 import time
 import typing
 from dataclasses import dataclass, field
@@ -69,9 +70,11 @@ from chalkcompute._function import (
     FunctionVersionInfo,
     RemoteFunction,
     SerializationFormat,
+    _assert_stripped_source_self_contained,
     _build_and_upload,
     _callable_input_names,
     _cleanup_temp_files,
+    _entrypoint_source,
     _extract_field_names,
     _is_chalkcompute_decorator,
     _local_install_entrypoint,
@@ -508,74 +511,78 @@ def _generate_class_handler_shim(
 
     method_caller = "_instance" if method_has_self else class_name
     if batched:
-        parts += [
-            "",
-            f"_is_coroutine = inspect.iscoroutinefunction({method_caller}.{method_name})",
-            "",
-            "",
-            "def handler(events, contexts):",
-            "    if _is_coroutine:",
-            "        import asyncio",
-            "        async def _run_all():",
-            "            results = []",
-            "            for event in events:",
-            "                columns = {k: v.to_pylist() for k, v in event.items()}",
-            "                n = len(next(iter(columns.values())))",
-            "                per_caller = []",
-            "                for i in range(n):",
-            "                    kwargs = {k: v[i] for k, v in columns.items()}",
-            f"                    per_caller.append(await {method_caller}.{method_name}(**kwargs))",
-            "                results.append(per_caller)",
-            "            return results",
-            "        return asyncio.run(_run_all())",
-            "    results = []",
-            "    for event in events:",
-            "        columns = {k: v.to_pylist() for k, v in event.items()}",
-            "        n = len(next(iter(columns.values())))",
-            "        per_caller = []",
-            "        for i in range(n):",
-            "            kwargs = {k: v[i] for k, v in columns.items()}",
-            f"            per_caller.append({method_caller}.{method_name}(**kwargs))",
-            "        results.append(per_caller)",
-            "    return results",
-        ]
+        parts.append(
+            textwrap.dedent(f"""\
+            _target = inspect.unwrap({method_caller}.{method_name})
+            _is_coroutine = inspect.iscoroutinefunction(_target)
+
+
+            def handler(events, contexts):
+                if _is_coroutine:
+                    import asyncio
+                    async def _run_all():
+                        results = []
+                        for event in events:
+                            columns = {{k: v.to_pylist() for k, v in event.items()}}
+                            n = len(next(iter(columns.values())))
+                            per_caller = []
+                            for i in range(n):
+                                kwargs = {{k: v[i] for k, v in columns.items()}}
+                                per_caller.append(await {method_caller}.{method_name}(**kwargs))
+                            results.append(per_caller)
+                        return results
+                    return asyncio.run(_run_all())
+                results = []
+                for event in events:
+                    columns = {{k: v.to_pylist() for k, v in event.items()}}
+                    n = len(next(iter(columns.values())))
+                    per_caller = []
+                    for i in range(n):
+                        kwargs = {{k: v[i] for k, v in columns.items()}}
+                        per_caller.append({method_caller}.{method_name}(**kwargs))
+                    results.append(per_caller)
+                return results
+            """).rstrip()
+        )
     else:
-        parts += [
-            "",
-            f"_is_sync_gen = inspect.isgeneratorfunction({method_caller}.{method_name})",
-            f"_is_async_gen = inspect.isasyncgenfunction({method_caller}.{method_name})",
-            f"_is_coroutine = inspect.iscoroutinefunction({method_caller}.{method_name}) and not _is_async_gen",
-            "",
-            "",
-            "def handler(event, context):",
-            "    n = len(next(iter(event.values())))",
-            "    if _is_sync_gen:",
-            "        def _gen():",
-            "            for i in range(n):",
-            "                kwargs = {k: v[i].as_py() for k, v in event.items()}",
-            f"                yield from {method_caller}.{method_name}(**kwargs)",
-            "        return _gen()",
-            "    elif _is_async_gen:",
-            "        async def _agen():",
-            "            for i in range(n):",
-            "                kwargs = {k: v[i].as_py() for k, v in event.items()}",
-            f"                async for item in {method_caller}.{method_name}(**kwargs):",
-            "                    yield item",
-            "        return _agen()",
-            "    elif _is_coroutine:",
-            "        import asyncio",
-            "        results = []",
-            "        for i in range(n):",
-            "            kwargs = {k: v[i].as_py() for k, v in event.items()}",
-            f"            results.append(asyncio.run({method_caller}.{method_name}(**kwargs)))",
-            "        return results",
-            "    else:",
-            "        results = []",
-            "        for i in range(n):",
-            "            kwargs = {k: v[i].as_py() for k, v in event.items()}",
-            f"            results.append({method_caller}.{method_name}(**kwargs))",
-            "        return results",
-        ]
+        parts.append(
+            textwrap.dedent(f"""\
+            _target = inspect.unwrap({method_caller}.{method_name})
+            _is_sync_gen = inspect.isgeneratorfunction(_target)
+            _is_async_gen = inspect.isasyncgenfunction(_target)
+            _is_coroutine = inspect.iscoroutinefunction(_target) and not _is_async_gen
+
+
+            def handler(event, context):
+                n = len(next(iter(event.values())))
+                if _is_sync_gen:
+                    def _gen():
+                        for i in range(n):
+                            kwargs = {{k: v[i].as_py() for k, v in event.items()}}
+                            yield from {method_caller}.{method_name}(**kwargs)
+                    return _gen()
+                elif _is_async_gen:
+                    async def _agen():
+                        for i in range(n):
+                            kwargs = {{k: v[i].as_py() for k, v in event.items()}}
+                            async for item in {method_caller}.{method_name}(**kwargs):
+                                yield item
+                    return _agen()
+                elif _is_coroutine:
+                    import asyncio
+                    results = []
+                    for i in range(n):
+                        kwargs = {{k: v[i].as_py() for k, v in event.items()}}
+                        results.append(asyncio.run({method_caller}.{method_name}(**kwargs)))
+                    return results
+                else:
+                    results = []
+                    for i in range(n):
+                        kwargs = {{k: v[i].as_py() for k, v in event.items()}}
+                        results.append({method_caller}.{method_name}(**kwargs))
+                    return results
+            """).rstrip()
+        )
 
     return "\n".join(parts) + "\n"
 
@@ -931,6 +938,15 @@ class RemoteClass:
         if _should_strip_source():
             source = _source_text_for_deploy(source_file)
             stripped = _strip_chalkcompute_from_class_source(source)
+            entrypoint_sources: list[str] = []
+            for method_info in self._methods.values():
+                entrypoint_sources.extend(_entrypoint_source(method_info.func))
+            _assert_stripped_source_self_contained(
+                source,
+                stripped,
+                label=f"Class {self._class_name!r}",
+                entrypoint_sources=entrypoint_sources,
+            )
             source_file = _write_temp(stripped)
             temp_files.append(source_file)
         elif not source_path.is_file():