1
import signal
2
import sys
3
from inspect import isawaitable, isclass
4
from types import FunctionType
5
 
6
import js
7
from workers import (
8
    DurableObject,
9
    WorkerEntrypoint,
10
    WorkflowEntrypoint,
11
    python_from_rpc,
12
    python_to_rpc,
13
)
14
 
15
from pyodide.code import relaxed_call
16
from pyodide.ffi import to_js
17
 
18
 
19
def getattr_no_get(cls, name):
20
    """Get attribute from class, don't run __get__"""
21
    for base in cls.__mro__:
22
        if name in base.__dict__:
23
            return base.__dict__[name]
24
    return None
25
 
26
 
27
def collect_methods(cls):
28
    """
29
    Iterates through the methods in `cls` and returns only public non-static/non-class methods
30
    defined on that class.
31
    """
32
    return sorted(
33
        name
34
        for name in dir(cls)
35
        if not name.startswith("_")
36
        and isinstance(getattr_no_get(cls, name), FunctionType)
37
    )
38
 
39
 
40
def collect_classes(user_mod, base_cls):
41
    """
42
    Iterates through the defined symbols in the input module. Returns any classes which extend
43
    `base_cls` (where `base_cls` is one of DurableObject, WorkerEntrypoint or WorkflowEntrypoint).
44
45
    This method returns a list of JS objects like [{"className": "MyClass", "methodNames": ["foo"]}].]
46
    """
47
    if hasattr(user_mod, "__all__"):
48
        keys = user_mod.__all__
49
    else:
50
        keys = (key for key in dir(user_mod) if not key.startswith("_"))
51
 
52
    exported_attrs = [getattr(user_mod, key) for key in keys]
53
 
54
    def filter(val):
55
        return isclass(val) and issubclass(val, base_cls) and val is not base_cls
56
 
57
    class_attrs = [attr for attr in exported_attrs if filter(attr)]
58
    result = [
59
        {"className": attr.__name__, "methodNames": collect_methods(attr)}
60
        for attr in class_attrs
61
    ]
62
    return to_js(result, dict_converter=js.Object.fromEntries)
63
 
64
 
65
def collect_entrypoint_classes(user_mod):
66
    return to_js(
67
        {
68
            "durableObjects": collect_classes(user_mod, DurableObject),
69
            "workerEntrypoints": collect_classes(user_mod, WorkerEntrypoint),
70
            "workflowEntrypoints": collect_classes(user_mod, WorkflowEntrypoint),
71
        },
72
        dict_converter=js.Object.fromEntries,
73
    )
74
 
75
 
76
async def wrapper_func(relaxed, inst, prop, *args, **kwargs):
77
    method = getattr(inst, prop)
78
 
79
    py_args = [python_from_rpc(arg) for arg in args]
80
    py_kwargs = {k: python_from_rpc(v) for k, v in kwargs.items()}
81
    result = (
82
        relaxed_call(method, *py_args, **py_kwargs)
83
        if relaxed
84
        else method(*py_args, **py_kwargs)
85
    )
86
 
87
    if isawaitable(result):
88
        return python_to_rpc(await result)
89
    else:
90
        return python_to_rpc(result)
91
 
92
 
93
class CpuLimitExceeded(BaseException):
94
    pass
95
 
96
 
97
def raise_cpu_limit_exceeded(signum, frame):
98
    raise CpuLimitExceeded("Python Worker exceeded CPU time limit")
99
 
100
 
101
if sys.platform == "emscripten":
102
    signal.signal(signal.SIGXCPU, raise_cpu_limit_exceeded)