pydiverse-common 0.3.2__py2.py3-none-any.whl

pydiverse/common/util/computation_tracing.py

@@ -0,0 +1,341 @@
+# Copyright (c) QuantCo and pydiverse contributors 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+import dis
+import inspect
+from enum import Enum
+from typing import Any
+
+
+class Operation(Enum):
+    OBJECT = 0
+    GETATTR = 10
+    SETATTR = 11
+    DELATTR = 12
+    GETITEM = 20
+    SETITEM = 21
+    DELITEM = 22
+    CALL = 50
+    GET = 60
+    BOOL = 70
+
+    def __repr__(self):
+        return self.name
+
+
+class ComputationTracer:
+    proxy_type: type["ComputationTracerProxy"]
+
+    def __init__(self):
+        self.trace = []
+        self.patcher = MonkeyPatcher()
+        self.did_exit = False
+        self.proxy_type = ComputationTracerProxy
+
+    def create_proxy(self, identifier=None):
+        return self._get_proxy((Operation.OBJECT, identifier))
+
+    def __enter__(self):
+        self._monkey_patch()
+        # clear trace already filled during patching (modules may issue calls during
+        # initialization)
+        self.trace = []
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.patcher.undo()
+        self.did_exit = True
+
+    def _get_proxy(self, computation: tuple):
+        idx = len(self.trace)
+        self._add_computation(computation)
+        return self.proxy_type(self, idx)
+
+    def _add_computation(self, computation: tuple):
+        if not self.did_exit:
+            from pydiverse.common.util.deep_map import deep_map
+
+            computation = deep_map(computation, self._computation_mapper)
+            self.trace.append(computation)
+        else:
+            raise RuntimeError(
+                "Can't modify ComputationTrace after exiting the context."
+            )
+
+    @staticmethod
+    def _computation_mapper(x):
+        if isinstance(x, ComputationTracerProxy):
+            return ComputationTraceRef(x)
+
+        if inspect.isfunction(x):
+            bytecode = dis.Bytecode(x)
+            return (
+                Operation.OBJECT,
+                "BYTECODE",
+                tuple((instr.opcode, instr.argval) for instr in bytecode),
+            )
+
+        return x
+
+    def _monkey_patch(self): ...
+
+    def trace_hash(self) -> str:
+        try:
+            from dask.base import tokenize
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "module dask is required to use computation_tracing."
+            ) from None
+
+        return tokenize(self.trace)
+
+
+class ComputationTracerProxy:
+    def __init__(self, tracer: ComputationTracer, identifier: int | str):
+        object.__setattr__(self, "_computation_tracer_", tracer)
+        object.__setattr__(self, "_computation_tracer_id_", identifier)
+
+    def __getattribute__(self, item):
+        if item in ("__class__", "__module__"):
+            return object.__getattribute__(self, item)
+
+        tracer = _get_tracer(self)
+        return tracer._get_proxy((Operation.GETATTR, self, item))
+
+    def __setattr__(self, key, value):
+        tracer = _get_tracer(self)
+        tracer._add_computation((Operation.SETATTR, self, key, value))
+
+    def __delattr__(self, key):
+        tracer = _get_tracer(self)
+        tracer._add_computation((Operation.DELATTR, self, key))
+
+    def __getitem__(self, item):
+        tracer = _get_tracer(self)
+        return tracer._get_proxy((Operation.GETITEM, self, item))
+
+    def __setitem__(self, key, value):
+        tracer = _get_tracer(self)
+        tracer._add_computation((Operation.SETITEM, self, key, value))
+
+    def __delitem__(self, key):
+        tracer = _get_tracer(self)
+        tracer._add_computation((Operation.DELITEM, self, key))
+
+    def __call__(self, *args, **kwargs):
+        tracer = _get_tracer(self)
+        return tracer._get_proxy((Operation.CALL, self, args, kwargs))
+
+    def __get__(self, instance, owner):
+        tracer = _get_tracer(self)
+        return tracer._get_proxy((Operation.GET, instance, self))
+
+    def __bool__(self):
+        tracer = _get_tracer(self)
+        return tracer._get_proxy((Operation.BOOL, self))
+
+    def __iter__(self):
+        raise RuntimeError("__iter__ is not supported by ComputationTracerProxy")
+
+    def __contains__(self, item):
+        raise RuntimeError("__contains__ is not supported by ComputationTracerProxy")
+
+    def __len__(self):
+        raise RuntimeError("__len__ is not supported by ComputationTracerProxy")
+
+
+def _get_tracer(proxy: ComputationTracerProxy) -> ComputationTracer:
+    return object.__getattribute__(proxy, "_computation_tracer_")
+
+
+__supported_dunder = {
+    "__add__",
+    "__radd__",
+    "__sub__",
+    "__rsub__",
+    "__mul__",
+    "__rmul__",
+    "__truediv__",
+    "__rtruediv__",
+    "__floordiv__",
+    "__rfloordiv__",
+    "__pow__",
+    "__rpow__",
+    "__mod__",
+    "__rmod__",
+    "__round__",
+    "__pos__",
+    "__neg__",
+    "__abs__",
+    "__and__",
+    "__rand__",
+    "__or__",
+    "__ror__",
+    "__xor__",
+    "__rxor__",
+    "__invert__",
+    "__lt__",
+    "__le__",
+    "__eq__",
+    "__ne__",
+    "__gt__",
+    "__ge__",
+    "__copy__",
+    "__deepcopy__",
+}
+
+
+def __create_dunder(name):
+    def dunder(self, *args):
+        return getattr(self, name)(self, *args)
+
+    return dunder
+
+
+for dunder_ in __supported_dunder:
+    setattr(ComputationTracerProxy, dunder_, __create_dunder(dunder_))
+
+
+class ComputationTraceRef:
+    __slots__ = ("id",)
+
+    def __init__(self, proxy: ComputationTracerProxy):
+        self.id = object.__getattribute__(proxy, "_computation_tracer_id_")
+
+    def __str__(self):
+        return f"ComputationTraceRef<{self.id}>"
+
+    def __repr__(self):
+        return f"ComputationTraceRef<{self.id}>"
+
+    def __dask_tokenize__(self):
+        return "ComputationTraceRef", self.id
+
+
+class MonkeyPatcher:
+    """Monkey Patching class inspired by pytest's MonkeyPatch class"""
+
+    def __init__(self):
+        self._setattr: list[tuple[object, str, Any]] = []
+
+    def patch_attr(self, obj: object, name: str, value: Any):
+        old_value = getattr(obj, name)
+
+        # avoid class descriptors like staticmethod / classmethod
+        if inspect.isclass(obj):
+            old_value = obj.__dict__[name]
+
+        setattr(obj, name, value)
+        self._setattr.append((obj, name, old_value))
+
+    def undo(self):
+        for obj, name, value in reversed(self._setattr):
+            setattr(obj, name, value)
+        self._setattr.clear()
+
+
+def fully_qualified_name(obj):
+    if type(obj).__name__ == "builtin_function_or_method":
+        if obj.__module__ is not None:
+            module = obj.__module__
+        else:
+            if inspect.isclass(obj.__self__):
+                module = obj.__self__.__module__
+            else:
+                module = obj.__self__.__class__.__module__
+        return f"{module}.{obj.__qualname__}"
+
+    if type(obj).__name__ == "function":
+        if hasattr(obj, "__wrapped__"):
+            qualname = obj.__wrapped__.__qualname__
+        else:
+            qualname = obj.__qualname__
+        return f"{obj.__module__}.{qualname}"
+
+    if type(obj).__name__ in (
+        "member_descriptor",
+        "method_descriptor",
+        "wrapper_descriptor",
+    ):
+        return f"{obj.__objclass__.__module__}.{obj.__qualname__}"
+
+    if type(obj).__name__ == "method":
+        if inspect.isclass(obj.__self__):
+            cls = obj.__self__.__qualname__
+        else:
+            cls = obj.__self__.__class__.__qualname__
+        return f"{obj.__self__.__module__}.{cls}.{obj.__name__}"
+
+    if type(obj).__name__ == "method-wrapper":
+        return f"{fully_qualified_name(obj.__self__)}.{obj.__name__}"
+
+    if type(obj).__name__ == "module":
+        return obj.__name__
+
+    if type(obj).__name__ == "property":
+        return f"{obj.fget.__module__}.{obj.fget.__qualname__}"
+
+    if inspect.isclass(obj):
+        return f"{obj.__module__}.{obj.__qualname__}"
+
+    return f"{obj.__class__.__module__}.{obj.__class__.__qualname__}"
+
+
+def patch(tracer: ComputationTracer, target: object, name: str):
+    if isinstance(target, ComputationTracerProxy):
+        return
+
+    val = getattr(target, name)
+
+    if isinstance(val, type):
+        return patch_type(tracer, target, name)
+    if inspect.isfunction(val):
+        return patch_function(tracer, target, name)
+    if isinstance(val, object):
+        return patch_value(tracer, target, name)
+
+    raise RuntimeError
+
+
+def patch_type(tracer: ComputationTracer, target: object, name: str):
+    val = getattr(target, name)
+    full_name = fully_qualified_name(val)
+
+    dunder_to_patch = (
+        "__getattr__",
+        "__setattr__",
+        "__delattr__",
+        "__getitem__",
+        "__setitem__",
+        "__delitem__",
+        "__call__",
+        "__repr__",
+        "__str__",
+    )
+
+    for key, _value in object.__getattribute__(val, "__dict__").items():
+        if key.startswith("__"):
+            if key in dunder_to_patch:
+                try:
+                    tracer.patcher.patch_attr(
+                        val, key, tracer.proxy_type(tracer, full_name + "." + key)
+                    )
+                except AttributeError:
+                    pass
+            continue
+        else:
+            tracer.patcher.patch_attr(
+                val, key, tracer.proxy_type(tracer, full_name + "." + key)
+            )
+
+    tracer.patcher.patch_attr(target, name, tracer.proxy_type(tracer, full_name))
+
+
+def patch_function(tracer: ComputationTracer, target: object, name: str):
+    val = getattr(target, name)
+    full_name = fully_qualified_name(val)
+    tracer.patcher.patch_attr(target, name, tracer.proxy_type(tracer, full_name))
+
+
+def patch_value(tracer: ComputationTracer, target: object, name: str):
+    full_name = fully_qualified_name(target) + "." + name
+    tracer.patcher.patch_attr(target, name, tracer.proxy_type(tracer, full_name))

pydiverse/common/util/deep_map.py

@@ -0,0 +1,100 @@
+# Copyright (c) QuantCo and pydiverse contributors 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Generic deep map or mutation operations.
+
+Heavily inspired by the builtin copy module of python:
+https://github.com/python/cpython/blob/main/Lib/copy.py
+"""
+
+from collections.abc import Callable
+
+from .computation_tracing import fully_qualified_name
+from .import_ import load_object
+
+_nil = []
+
+
+def deep_map(x, fn: Callable, memo=None):
+    if memo is None:
+        memo = {}
+
+    d = id(x)
+    y = memo.get(d, _nil)
+    if y is not _nil:
+        return y
+
+    cls = type(x)
+
+    if cls == list:  # noqa: E721
+        y = _deep_map_list(x, fn, memo)
+    elif cls == tuple:  # noqa: E721
+        y = _deep_map_tuple(x, fn, memo)
+    elif cls == dict:  # noqa: E721
+        y = _deep_map_dict(x, fn, memo)
+    elif hasattr(cls, "__dataclass_fields__"):
+        # reconstruct data classes
+        y = load_object(
+            {
+                "class": fully_qualified_name(cls),
+                "args": _deep_map_dict(x.__dict__, fn, memo),
+            }
+        )
+    else:
+        y = fn(x)
+
+    # If is its own copy, don't memoize.
+    if y is not x:
+        memo[d] = y
+        _keep_alive(x, memo)  # Make sure x lives at least as long as d
+
+    return y
+
+
+def _deep_map_list(x, fn, memo):
+    y = []
+    append = y.append
+    for a in x:
+        append(deep_map(a, fn, memo))
+    return fn(y)
+
+
+def _deep_map_tuple(x, fn, memo):
+    y = [deep_map(a, fn, memo) for a in x]
+    # We're not going to put the tuple in the memo, but it's still important we
+    # check for it, in case the tuple contains recursive mutable structures.
+    try:
+        return memo[id(x)]
+    except KeyError:
+        pass
+    for k, j in zip(x, y, strict=False):
+        if k is not j:
+            y = tuple(y)
+            break
+    else:
+        y = x
+    return fn(y)
+
+
+def _deep_map_dict(x, fn, memo):
+    y = {}
+    memo[id(x)] = y
+    for key, value in x.items():
+        y[deep_map(key, fn, memo)] = deep_map(value, fn, memo)
+    return fn(y)
+
+
+def _keep_alive(x, memo):
+    """Keeps a reference to the object x in the memo.
+    Because we remember objects by their id, we have
+    to assure that possibly temporary objects are kept
+    alive by referencing them.
+    We store a reference at the id of the memo, which should
+    normally not be used unless someone tries to deepcopy
+    the memo itself...
+    """
+    try:
+        memo[id(memo)].append(x)
+    except KeyError:
+        # aha, this is the first one :-)
+        memo[id(memo)] = [x]

pydiverse/common/util/deep_merge.py

@@ -0,0 +1,55 @@
+# Copyright (c) QuantCo and pydiverse contributors 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Generic deep update function for nested dictionaries.
+
+Seems to be solved already in various ways (do we like an extra dependency for pydantic.deep_update?)
+https://stackoverflow.com/questions/3232943/update-value-of-a-nested-dictionary-of-varying-depth
+But for snippets, license restrictions exist:
+https://www.ictrecht.nl/en/blog/what-is-the-license-status-of-stackoverflow-code-snippets
+"""  # noqa: E501
+
+from collections.abc import Iterable, Mapping
+
+from box import Box
+
+
+def deep_merge(x, y, check_enum=False):
+    if type(x) != type(y) and not (isinstance(x, Mapping) and isinstance(y, Mapping)):  # noqa: E721
+        raise TypeError(
+            f"deep_merge failed due to type mismatch '{x}' (type: {type(x)}) vs. '{y}'"
+            f" (type: {type(y)})"
+        )
+
+    if isinstance(x, Box):
+        z = Box(_deep_merge_dict(x, y), frozen_box=True)
+    elif isinstance(x, Mapping):
+        z = _deep_merge_dict(x, y)
+    elif isinstance(x, Iterable) and not isinstance(x, str):
+        z = _deep_merge_iterable(x, y)
+    else:
+        z = y  # update
+
+    return z
+
+
+def _deep_merge_iterable(x: Iterable, y: Iterable):
+    # Merging lists is not trivial.
+    # There are a few different strategies: replace, unique, append, intersection, ...
+    return y
+    # return [*x, *y]
+    # return [deep_merge(a, b) for a, b in zip(x, y)]
+
+
+def _deep_merge_dict(x: Mapping, y: Mapping):
+    z = dict(x)
+    for key in x:
+        if key in y:
+            if y[key] is None:
+                # this is a special case but we have no other way in yaml to express
+                # the deletion of fields from a dictionary in an override config
+                del z[key]
+            else:
+                z[key] = deep_merge(x[key], y[key])
+    z.update({key: value for key, value in y.items() if key not in z})
+    return z

pydiverse/common/util/disposable.py

@@ -0,0 +1,28 @@
+# Copyright (c) QuantCo and pydiverse contributors 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+from ..errors import DisposedError
+
+
+class Disposable:
+    def __getattribute__(self, name):
+        try:
+            object.__getattribute__(self, "_Disposable__disposed")
+            obj_type = object.__getattribute__(self, "__class__")
+            raise DisposedError(f"Object of type {obj_type} has already been disposed.")
+        except AttributeError:
+            pass
+
+        return object.__getattribute__(self, name)
+
+    def __setattr__(self, key, value):
+        try:
+            object.__getattribute__(self, "_Disposable__disposed")
+            obj_type = object.__getattribute__(self, "__class__")
+            raise DisposedError(f"Object of type {obj_type} has already been disposed.")
+        except AttributeError:
+            pass
+
+        return object.__setattr__(self, key, value)
+
+    def dispose(self):
+        object.__setattr__(self, "_Disposable__disposed", True)

pydiverse/common/util/hashing.py

@@ -0,0 +1,32 @@
+# Copyright (c) QuantCo and pydiverse contributors 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+import base64
+import hashlib
+
+
+def stable_hash(*args: str) -> str:
+    """Compute a hash over a set of strings
+
+    :param args: Some strings from which to compute the cache key
+    :return: A sha256 base32 digest, trimmed to 20 char length
+    """
+
+    combined_hash = hashlib.sha256(b"PYDIVERSE")
+    for arg in args:
+        arg_bytes = str(arg).encode("utf8")
+        arg_bytes_len = len(arg_bytes).to_bytes(length=8, byteorder="big")
+
+        combined_hash.update(arg_bytes_len)
+        combined_hash.update(arg_bytes)
+
+    # Only take first 20 characters of base32 digest (100 bits). This
+    # provides 50 bits of collision resistance, which is more than enough.
+    # To illustrate: If you were to generate 1k hashes per second,
+    # you still would have to wait over 800k years until you encounter
+    # a collision.
+
+    # NOTE: Can't use base64 because it contains lower and upper case
+    #       letters; identifiers in pipedag are all lowercase
+    hash_digest = combined_hash.digest()
+    hash_str = base64.b32encode(hash_digest).decode("ascii").lower()
+    return hash_str[:20]

pydiverse/common/util/import_.py

@@ -0,0 +1,135 @@
+# Copyright (c) QuantCo and pydiverse contributors 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+import builtins
+import importlib
+import os
+from collections.abc import Collection
+from typing import Any
+
+_allowed_getattr = [
+    "__class__",
+    "__doc__",
+    "__name__",
+    "__qualname__",
+    "__module__",
+]
+
+
+def requires(requirements: Any | list, exception: BaseException | type[BaseException]):
+    """Class decorator for handling optional imports.
+
+    If any of the requirements are falsy, this decorator prevents the class
+    from being instantiated and any class attributes from being accessed,
+    and raises the provided exception instead.
+    """
+
+    if not isinstance(requirements, list | tuple):
+        requirements = (requirements,)
+
+    def decorator(cls):
+        if all(requirements):
+            return cls
+
+        # Modify class to raise exception
+        class RaiserMeta(type):
+            def __getattribute__(self, x):
+                # While building the documentation, we set the SPHINX_BUILD env
+                # variable. This allows us to properly generate the documentation
+                # without raising exceptions.
+                if os.environ.get("SPHINX_BUILD"):
+                    return getattr(cls, x)
+                if x in _allowed_getattr:
+                    return getattr(cls, x)
+                raise exception
+
+        def raiser(*args, **kwargs):
+            raise exception
+
+        __name = str(cls.__name__)
+        __bases = ()
+        __dict = {
+            "__metaclass__": RaiserMeta,
+            "__wrapped__": cls,
+            "__new__": raiser,
+        }
+
+        return RaiserMeta(__name, __bases, __dict)
+
+    return decorator
+
+
+def import_object(import_path: str):
+    """Loads a class given an import path
+
+    >>> # An import statement like this
+    >>> from pandas import DataFrame
+    >>> # can be expressed as follows:
+    >>> import_object("pandas.DataFrame")
+    """
+
+    parts = [part for part in import_path.split(".") if part]
+    module, n = None, 0
+
+    while n < len(parts):
+        try:
+            module = importlib.import_module(".".join(parts[: n + 1]))
+            n = n + 1
+        except ImportError:
+            break
+
+    obj = module or builtins
+    for part in parts[n:]:
+        obj = getattr(obj, part)
+
+    return obj
+
+
+def load_object(config_dict: dict, move_keys_into_args: Collection[str] | None = None):
+    """Instantiates an instance of an object given
+
+    The import path (module.Class) should be specified as the "class" value
+    of the dict. The args section of the dict get used as the instance config.
+
+    If the class defines a `_init_conf_` function, it gets called using the
+    config values, otherwise they just get passed to the class initializer.
+
+    Additionally, any values of `config_dict` whose associated keys are in
+    `move_keys_into_args`, get also passed as an argument to the initializer.
+    ::
+
+        # module.Class(argument="value")
+        load_object({
+            "class": "module.Class",
+            "args": {
+                "argument": "value",
+            },
+        })
+    """
+
+    if "class" not in config_dict:
+        raise RuntimeError(
+            "Attribute 'class' is missing in configuration "
+            "section that supports multiple backends\n"
+            f"config section: {config_dict}"
+        )
+    if isinstance(config_dict["class"], type):
+        # it may be useful in tests to just pass in a dynamically created class
+        cls = config_dict["class"]
+    else:
+        cls = import_object(config_dict["class"])
+
+    args = config_dict.get("args", {}) or {}
+    if not isinstance(args, dict):
+        raise TypeError(
+            f"Invalid type for args section: {type(args)}\n"
+            f"config section: {config_dict}"
+        )
+
+    if move_keys_into_args:
+        args = args | {k: v for k, v in config_dict.items() if k in move_keys_into_args}
+
+    try:
+        init_conf = cls._init_conf_
+        return init_conf(args)
+    except AttributeError:
+        return cls(**args)