coco-framework 0.1.0__py3-none-any.whl

coco_graph/_internal/_fields.py

@@ -0,0 +1,213 @@
+from __future__ import annotations
+
+import dataclasses
+import types
+import weakref
+from collections.abc import Generator, Sequence
+from typing import Annotated, Any, Optional, Union, get_origin, get_type_hints
+
+from pydantic import BaseModel
+from typing_extensions import NotRequired, ReadOnly, Required
+
+from coco_graph._internal._typing import MISSING
+
+
+def _is_optional_type(type_: Any) -> bool:
+    """Check if a type is Optional."""
+
+    # Handle new union syntax (PEP 604): str | None
+    if isinstance(type_, types.UnionType):
+        return any(
+            arg is type(None) or _is_optional_type(arg) for arg in type_.__args__
+        )
+
+    if hasattr(type_, "__origin__") and hasattr(type_, "__args__"):
+        origin = get_origin(type_)
+        if origin is Optional:
+            return True
+        if origin is Union:
+            return any(
+                arg is type(None) or _is_optional_type(arg) for arg in type_.__args__
+            )
+        if origin is Annotated:
+            return _is_optional_type(type_.__args__[0])
+        return origin is None
+    if hasattr(type_, "__bound__") and type_.__bound__ is not None:
+        return _is_optional_type(type_.__bound__)
+    return type_ is None
+
+
+def _is_required_type(type_: Any) -> bool | None:
+    """Check if an annotation is marked as Required/NotRequired.
+
+    Returns:
+        - True if required
+        - False if not required
+        - None if not annotated with either
+    """
+    origin = get_origin(type_)
+    if origin is Required:
+        return True
+    if origin is NotRequired:
+        return False
+    if origin is Annotated or getattr(origin, "__args__", None):
+        # See https://typing.readthedocs.io/en/latest/spec/typeddict.html#interaction-with-annotated
+        return _is_required_type(type_.__args__[0])
+    return None
+
+
+def _is_readonly_type(type_: Any) -> bool:
+    """Check if an annotation is marked as ReadOnly.
+
+    Returns:
+        - True if is read only
+        - False if not read only
+    """
+
+    # See: https://typing.readthedocs.io/en/latest/spec/typeddict.html#typing-readonly-type-qualifier
+    origin = get_origin(type_)
+    if origin is Annotated:
+        return _is_readonly_type(type_.__args__[0])
+    if origin is ReadOnly:
+        return True
+    return False
+
+
+_DEFAULT_KEYS: frozenset[str] = frozenset()
+
+
+def get_field_default(name: str, type_: Any, schema: type[Any]) -> Any:
+    """Determine the default value for a field in a state schema.
+
+    This is based on:
+        If TypedDict:
+            - Required/NotRequired
+            - total=False -> everything optional
+        - Type annotation (Optional/Union[None])
+    """
+    optional_keys = getattr(schema, "__optional_keys__", _DEFAULT_KEYS)
+    irq = _is_required_type(type_)
+    if name in optional_keys:
+        # Either total=False or explicit NotRequired.
+        # No type annotation trumps this.
+        if irq:
+            # Unless it's earlier versions of python & explicit Required
+            return ...
+        return None
+    if irq is not None:
+        if irq:
+            # Handle Required[<type>]
+            # (we already handled NotRequired and total=False)
+            return ...
+        # Handle NotRequired[<type>] for earlier versions of python
+        return None
+    if dataclasses.is_dataclass(schema):
+        field_info = next(
+            (f for f in dataclasses.fields(schema) if f.name == name), None
+        )
+        if field_info:
+            if (
+                field_info.default is not dataclasses.MISSING
+                and field_info.default is not ...
+            ):
+                return field_info.default
+            elif field_info.default_factory is not dataclasses.MISSING:
+                return field_info.default_factory()
+    # Note, we ignore ReadOnly attributes,
+    # as they don't make much sense. (we don't care if you mutate the state in your node)
+    # and mutating state in your node has no effect on our graph state.
+    # Base case is the annotation
+    if _is_optional_type(type_):
+        return None
+    return ...
+
+
+def get_enhanced_type_hints(
+    type: type[Any],
+) -> Generator[tuple[str, Any, Any, str | None], None, None]:
+    """Attempt to extract default values and descriptions from provided type, used for config schema."""
+    for name, typ in get_type_hints(type).items():
+        default = None
+        description = None
+
+        # Pydantic models
+        try:
+            if hasattr(type, "model_fields") and name in type.model_fields:
+                field = type.model_fields[name]
+
+                if hasattr(field, "description") and field.description is not None:
+                    description = field.description
+
+                if hasattr(field, "default") and field.default is not None:
+                    default = field.default
+                    if (
+                        hasattr(default, "__class__")
+                        and getattr(default.__class__, "__name__", "")
+                        == "PydanticUndefinedType"
+                    ):
+                        default = None
+
+        except (AttributeError, KeyError, TypeError):
+            pass
+
+        # TypedDict, dataclass
+        try:
+            if hasattr(type, "__dict__"):
+                type_dict = getattr(type, "__dict__")
+
+                if name in type_dict:
+                    default = type_dict[name]
+        except (AttributeError, KeyError, TypeError):
+            pass
+
+        yield name, typ, default, description
+
+
+def get_update_as_tuples(input: Any, keys: Sequence[str]) -> list[tuple[str, Any]]:
+    """Get Pydantic state update as a list of (key, value) tuples."""
+    if isinstance(input, BaseModel):
+        keep = input.model_fields_set
+        defaults = {k: v.default for k, v in type(input).model_fields.items()}
+    else:
+        keep = None
+        defaults = {}
+
+    # NOTE: This behavior for Pydantic is somewhat inelegant,
+    # but we keep around for backwards compatibility
+    # if input is a Pydantic model, only update values
+    # that are different from the default values or in the keep set
+    return [
+        (k, value)
+        for k in keys
+        if (value := getattr(input, k, MISSING)) is not MISSING
+        and (
+            value is not None
+            or defaults.get(k, MISSING) is not None
+            or (keep is not None and k in keep)
+        )
+    ]
+
+
+ANNOTATED_KEYS_CACHE: weakref.WeakKeyDictionary[type[Any], tuple[str, ...]] = (
+    weakref.WeakKeyDictionary()
+)
+
+
+def get_cached_annotated_keys(obj: type[Any]) -> tuple[str, ...]:
+    """Return cached annotated keys for a Python class."""
+    if obj in ANNOTATED_KEYS_CACHE:
+        return ANNOTATED_KEYS_CACHE[obj]
+    if isinstance(obj, type):
+        keys: list[str] = []
+        for base in reversed(obj.__mro__):
+            ann = base.__dict__.get("__annotations__")
+            # In Python 3.14+, Pydantic models use descriptors for __annotations__
+            # so we need to fall back to getattr if __dict__.get returns None
+            if ann is None:
+                ann = getattr(base, "__annotations__", None)
+            if ann is None or isinstance(ann, types.GetSetDescriptorType):
+                continue
+            keys.extend(ann.keys())
+        return ANNOTATED_KEYS_CACHE.setdefault(obj, tuple(keys))
+    else:
+        raise TypeError(f"Expected a type, got {type(obj)}. ")

coco_graph/_internal/_future.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import contextvars
+import inspect
+import sys
+import types
+from collections.abc import Awaitable, Coroutine, Generator
+from typing import TypeVar, cast
+
+T = TypeVar("T")
+AnyFuture = asyncio.Future | concurrent.futures.Future
+
+CONTEXT_NOT_SUPPORTED = sys.version_info < (3, 11)
+EAGER_NOT_SUPPORTED = sys.version_info < (3, 12)
+
+
+def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop:
+    # Tries to call Future.get_loop() if it's available.
+    # Otherwise fallbacks to using the old '_loop' property.
+    try:
+        get_loop = fut.get_loop
+    except AttributeError:
+        pass
+    else:
+        return get_loop()
+    return fut._loop
+
+
+def _convert_future_exc(exc: BaseException) -> BaseException:
+    exc_class = type(exc)
+    if exc_class is concurrent.futures.CancelledError:
+        return asyncio.CancelledError(*exc.args)
+    elif exc_class is concurrent.futures.TimeoutError:
+        return asyncio.TimeoutError(*exc.args)
+    elif exc_class is concurrent.futures.InvalidStateError:
+        return asyncio.InvalidStateError(*exc.args)
+    else:
+        return exc
+
+
+def _set_concurrent_future_state(
+    concurrent: concurrent.futures.Future,
+    source: AnyFuture,
+) -> None:
+    """Copy state from a future to a concurrent.futures.Future."""
+    assert source.done()
+    if source.cancelled():
+        concurrent.cancel()
+    if not concurrent.set_running_or_notify_cancel():
+        return
+    exception = source.exception()
+    if exception is not None:
+        concurrent.set_exception(_convert_future_exc(exception))
+    else:
+        result = source.result()
+        concurrent.set_result(result)
+
+
+def _copy_future_state(source: AnyFuture, dest: asyncio.Future) -> None:
+    """Internal helper to copy state from another Future.
+
+    The other Future may be a concurrent.futures.Future.
+    """
+    if dest.done():
+        return
+    assert source.done()
+    if dest.cancelled():
+        return
+    if source.cancelled():
+        dest.cancel()
+    else:
+        exception = source.exception()
+        if exception is not None:
+            dest.set_exception(_convert_future_exc(exception))
+        else:
+            result = source.result()
+            dest.set_result(result)
+
+
+def _chain_future(source: AnyFuture, destination: AnyFuture) -> None:
+    """Chain two futures so that when one completes, so does the other.
+
+    The result (or exception) of source will be copied to destination.
+    If destination is cancelled, source gets cancelled too.
+    Compatible with both asyncio.Future and concurrent.futures.Future.
+    """
+    if not asyncio.isfuture(source) and not isinstance(
+        source, concurrent.futures.Future
+    ):
+        raise TypeError("A future is required for source argument")
+    if not asyncio.isfuture(destination) and not isinstance(
+        destination, concurrent.futures.Future
+    ):
+        raise TypeError("A future is required for destination argument")
+    source_loop = _get_loop(source) if asyncio.isfuture(source) else None
+    dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None
+
+    def _set_state(future: AnyFuture, other: AnyFuture) -> None:
+        if asyncio.isfuture(future):
+            _copy_future_state(other, future)
+        else:
+            _set_concurrent_future_state(future, other)
+
+    def _call_check_cancel(destination: AnyFuture) -> None:
+        if destination.cancelled():
+            if source_loop is None or source_loop is dest_loop:
+                source.cancel()
+            else:
+                source_loop.call_soon_threadsafe(source.cancel)
+
+    def _call_set_state(source: AnyFuture) -> None:
+        if destination.cancelled() and dest_loop is not None and dest_loop.is_closed():
+            return
+        if dest_loop is None or dest_loop is source_loop:
+            _set_state(destination, source)
+        else:
+            if dest_loop.is_closed():
+                return
+            dest_loop.call_soon_threadsafe(_set_state, destination, source)
+
+    destination.add_done_callback(_call_check_cancel)
+    source.add_done_callback(_call_set_state)
+
+
+def chain_future(source: AnyFuture, destination: AnyFuture) -> AnyFuture:
+    # adapted from asyncio.run_coroutine_threadsafe
+    try:
+        _chain_future(source, destination)
+        return destination
+    except (SystemExit, KeyboardInterrupt):
+        raise
+    except BaseException as exc:
+        if isinstance(destination, concurrent.futures.Future):
+            if destination.set_running_or_notify_cancel():
+                destination.set_exception(exc)
+        else:
+            destination.set_exception(exc)
+        raise
+
+
+def _ensure_future(
+    coro_or_future: Coroutine[None, None, T] | Awaitable[T],
+    *,
+    loop: asyncio.AbstractEventLoop,
+    name: str | None = None,
+    context: contextvars.Context | None = None,
+    lazy: bool = True,
+) -> asyncio.Task[T]:
+    called_wrap_awaitable = False
+    if not asyncio.iscoroutine(coro_or_future):
+        if inspect.isawaitable(coro_or_future):
+            coro_or_future = cast(
+                Coroutine[None, None, T], _wrap_awaitable(coro_or_future)
+            )
+            called_wrap_awaitable = True
+        else:
+            raise TypeError(
+                "An asyncio.Future, a coroutine or an awaitable is required."
+                f" Got {type(coro_or_future).__name__} instead."
+            )
+
+    try:
+        if CONTEXT_NOT_SUPPORTED:
+            return loop.create_task(coro_or_future, name=name)
+        elif EAGER_NOT_SUPPORTED or lazy:
+            return loop.create_task(coro_or_future, name=name, context=context)
+        else:
+            return asyncio.eager_task_factory(
+                loop, coro_or_future, name=name, context=context
+            )
+    except RuntimeError:
+        if not called_wrap_awaitable:
+            coro_or_future.close()
+        raise
+
+
+@types.coroutine
+def _wrap_awaitable(awaitable: Awaitable[T]) -> Generator[None, None, T]:
+    """Helper for asyncio.ensure_future().
+
+    Wraps awaitable (an object with __await__) into a coroutine
+    that will later be wrapped in a Task by ensure_future().
+    """
+    return (yield from awaitable.__await__())
+
+
+def run_coroutine_threadsafe(
+    coro: Coroutine[None, None, T],
+    loop: asyncio.AbstractEventLoop,
+    *,
+    lazy: bool,
+    name: str | None = None,
+    context: contextvars.Context | None = None,
+) -> asyncio.Future[T]:
+    """Submit a coroutine object to a given event loop.
+
+    Return an asyncio.Future to access the result.
+    """
+
+    if asyncio._get_running_loop() is loop:
+        return _ensure_future(coro, loop=loop, name=name, context=context, lazy=lazy)
+    else:
+        future: asyncio.Future[T] = asyncio.Future(loop=loop)
+
+        def callback() -> None:
+            try:
+                chain_future(
+                    _ensure_future(coro, loop=loop, name=name, context=context),
+                    future,
+                )
+            except (SystemExit, KeyboardInterrupt):
+                raise
+            except BaseException as exc:
+                future.set_exception(exc)
+                raise
+
+        loop.call_soon_threadsafe(callback, context=context)
+        return future

coco_graph/_internal/_pydantic.py

@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+import sys
+import typing
+import warnings
+from contextlib import nullcontext
+from dataclasses import is_dataclass
+from functools import lru_cache
+from typing import (
+    Any,
+    cast,
+    overload,
+)
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    RootModel,
+)
+from pydantic import (
+    create_model as _create_model_base,
+)
+from pydantic.fields import FieldInfo
+from pydantic.json_schema import (
+    DEFAULT_REF_TEMPLATE,
+    GenerateJsonSchema,
+    JsonSchemaMode,
+)
+from typing_extensions import TypedDict
+
+
+@overload
+def get_fields(model: type[BaseModel]) -> dict[str, FieldInfo]: ...
+
+
+@overload
+def get_fields(model: BaseModel) -> dict[str, FieldInfo]: ...
+
+
+def get_fields(
+    model: type[BaseModel] | BaseModel,
+) -> dict[str, FieldInfo]:
+    """Get the field names of a Pydantic model."""
+    if hasattr(model, "model_fields"):
+        return model.model_fields
+
+    if hasattr(model, "__fields__"):
+        return model.__fields__
+    msg = f"Expected a Pydantic model. Got {type(model)}"
+    raise TypeError(msg)
+
+
+_SchemaConfig = ConfigDict(
+    arbitrary_types_allowed=True, frozen=True, protected_namespaces=()
+)
+
+NO_DEFAULT = object()
+
+
+def _create_root_model(
+    name: str,
+    type_: Any,
+    module_name: str | None = None,
+    default_: object = NO_DEFAULT,
+) -> type[BaseModel]:
+    """Create a base class."""
+
+    def schema(
+        cls: type[BaseModel],
+        by_alias: bool = True,
+        ref_template: str = DEFAULT_REF_TEMPLATE,
+    ) -> dict[str, Any]:
+        # Complains about schema not being defined in superclass
+        schema_ = super(cls, cls).schema(  # type: ignore[misc]
+            by_alias=by_alias, ref_template=ref_template
+        )
+        schema_["title"] = name
+        return schema_
+
+    def model_json_schema(
+        cls: type[BaseModel],
+        by_alias: bool = True,
+        ref_template: str = DEFAULT_REF_TEMPLATE,
+        schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema,
+        mode: JsonSchemaMode = "validation",
+    ) -> dict[str, Any]:
+        # Complains about model_json_schema not being defined in superclass
+        schema_ = super(cls, cls).model_json_schema(  # type: ignore[misc]
+            by_alias=by_alias,
+            ref_template=ref_template,
+            schema_generator=schema_generator,
+            mode=mode,
+        )
+        schema_["title"] = name
+        return schema_
+
+    base_class_attributes = {
+        "__annotations__": {"root": type_},
+        "model_config": ConfigDict(arbitrary_types_allowed=True),
+        "schema": classmethod(schema),
+        "model_json_schema": classmethod(model_json_schema),
+        "__module__": module_name or "coco_core.runnables.utils",
+    }
+
+    if default_ is not NO_DEFAULT:
+        base_class_attributes["root"] = default_
+    with warnings.catch_warnings():
+        custom_root_type = type(name, (RootModel,), base_class_attributes)
+    return cast("type[BaseModel]", custom_root_type)
+
+
+@lru_cache(maxsize=256)
+def _create_root_model_cached(
+    model_name: str,
+    type_: Any,
+    *,
+    module_name: str | None = None,
+    default_: object = NO_DEFAULT,
+) -> type[BaseModel]:
+    return _create_root_model(
+        model_name, type_, default_=default_, module_name=module_name
+    )
+
+
+@lru_cache(maxsize=256)
+def _create_model_cached(
+    model_name: str,
+    /,
+    **field_definitions: Any,
+) -> type[BaseModel]:
+    return _create_model_base(
+        model_name,
+        __config__=_SchemaConfig,
+        **_remap_field_definitions(field_definitions),
+    )
+
+
+# Reserved names should capture all the `public` names / methods that are
+# used by BaseModel internally. This will keep the reserved names up-to-date.
+# For reference, the reserved names are:
+# "construct", "copy", "dict", "from_orm", "json", "parse_file", "parse_obj",
+# "parse_raw", "schema", "schema_json", "update_forward_refs", "validate",
+# "model_computed_fields", "model_config", "model_construct", "model_copy",
+# "model_dump", "model_dump_json", "model_extra", "model_fields",
+# "model_fields_set", "model_json_schema", "model_parametrized_name",
+# "model_post_init", "model_rebuild", "model_validate", "model_validate_json",
+# "model_validate_strings"
+_RESERVED_NAMES = {key for key in dir(BaseModel) if not key.startswith("_")}
+
+
+def _remap_field_definitions(field_definitions: dict[str, Any]) -> dict[str, Any]:
+    """This remaps fields to avoid colliding with internal pydantic fields."""
+
+    remapped = {}
+    for key, value in field_definitions.items():
+        if key.startswith("_") or key in _RESERVED_NAMES:
+            # Let's add a prefix to avoid colliding with internal pydantic fields
+            if isinstance(value, FieldInfo):
+                msg = (
+                    f"Remapping for fields starting with '_' or fields with a name "
+                    f"matching a reserved name {_RESERVED_NAMES} is not supported if "
+                    f" the field is a pydantic Field instance. Got {key}."
+                )
+                raise NotImplementedError(msg)
+            type_, default_ = value
+            remapped[f"private_{key}"] = (
+                type_,
+                Field(
+                    default=default_,
+                    alias=key,
+                    serialization_alias=key,
+                    title=key.lstrip("_").replace("_", " ").title(),
+                ),
+            )
+        else:
+            remapped[key] = value
+    return remapped
+
+
+def create_model(
+    model_name: str,
+    *,
+    field_definitions: dict[str, Any] | None = None,
+    root: Any | None = None,
+) -> type[BaseModel]:
+    """Create a pydantic model with the given field definitions.
+
+    Attention:
+        Please do not use outside of coco packages. This API
+        is subject to change at any time.
+
+    Args:
+        model_name: The name of the model.
+        module_name: The name of the module where the model is defined.
+            This is used by Pydantic to resolve any forward references.
+        field_definitions: The field definitions for the model.
+        root: Type for a root model (RootModel)
+
+    Returns:
+        Type[BaseModel]: The created model.
+    """
+    field_definitions = field_definitions or {}
+
+    if root:
+        if field_definitions:
+            msg = (
+                "When specifying __root__ no other "
+                f"fields should be provided. Got {field_definitions}"
+            )
+            raise NotImplementedError(msg)
+
+        if isinstance(root, tuple):
+            kwargs = {"type_": root[0], "default_": root[1]}
+        else:
+            kwargs = {"type_": root}
+
+        try:
+            named_root_model = _create_root_model_cached(model_name, **kwargs)
+        except TypeError:
+            # something in the arguments into _create_root_model_cached is not hashable
+            named_root_model = _create_root_model(
+                model_name,
+                **kwargs,
+            )
+        return named_root_model
+
+    # No root, just field definitions
+    names = set(field_definitions.keys())
+
+    capture_warnings = False
+
+    for name in names:
+        # Also if any non-reserved name is used (e.g., model_id or model_name)
+        if name.startswith("model"):
+            capture_warnings = True
+
+    with warnings.catch_warnings() if capture_warnings else nullcontext():
+        if capture_warnings:
+            warnings.filterwarnings(action="ignore")
+        try:
+            return _create_model_cached(model_name, **field_definitions)
+        except TypeError:
+            # something in field definitions is not hashable
+            return _create_model_base(
+                model_name,
+                __config__=_SchemaConfig,
+                **_remap_field_definitions(field_definitions),
+            )
+
+
+def is_supported_by_pydantic(type_: Any) -> bool:
+    """Check if a given "complex" type is supported by pydantic.
+
+    This will return False for primitive types like int, str, etc.
+
+    The check is meant for container types like dataclasses, TypedDicts, etc.
+    """
+    if is_dataclass(type_):
+        return True
+
+    if isinstance(type_, type) and issubclass(type_, BaseModel):
+        return True
+
+    if hasattr(type_, "__orig_bases__"):
+        for base in type_.__orig_bases__:
+            if base is TypedDict:
+                return True
+            elif base is typing.TypedDict:  # noqa: TID251
+                # ignoring TID251 since it's OK to use typing.TypedDict in this case.
+                # Pydantic supports typing.TypedDict from Python 3.12
+                # For older versions, only typing_extensions.TypedDict is supported.
+                if sys.version_info >= (3, 12):
+                    return True
+    return False