diff --git a/docs/collections.md b/docs/collections.md new file mode 100644 index 0000000..9e31ca1 --- /dev/null +++ b/docs/collections.md @@ -0,0 +1,34 @@ +# Collections + +Requires the engine changes in [dagger/dagger#14221](https://github.com/dagger/dagger/pull/14221). + +A collection has stored keys and a function that returns one item for a key. +The engine supplies `keys`, `get`, `list`, and `subset`. Other exposed +functions appear under `batch`. + +```python +import dagger +from dagger import collection, delta, field, function, get, keys, object_type + +@object_type +class Item: + name: str = field() + +@collection +@object_type +class Items: + paths: list[str] = keys(default=list) + selection: dagger.CollectionDelta | None = delta() + + @function + @get + def item(self, key: str) -> Item: + return Item(name=key) +``` + +The markers are part of the module description. Both the shared entrypoint and +the generated static entrypoint pass them to the engine. + +The engine fills the optional delta field before a module call. It compares the +current keys with the original keys. Copies preserve the internal base state. +A new object starts a new base. The internal state is not an exposed field. diff --git a/entrypoint/main.dang b/entrypoint/main.dang index 86aa515..567f38d 100644 --- a/entrypoint/main.dang +++ b/entrypoint/main.dang @@ -109,17 +109,25 @@ type Entrypoint implements ModuleEntrypoint { deprecated: optString(raw, "deprecated"), ) } - let withFields = raw.field(["fields"]).asArray.{{contents}}.reduce(base) { acc, rawField => + let collection = if (raw.field(["collection"]).asBoolean) { base.withCollection } else { base } + let withFields = raw.field(["fields"]).asArray.{{contents}}.reduce(collection) { acc, rawField => let f = json.withContents(rawField.contents) - acc.withField( + let field = acc.withField( f.field(["name"]).asString, typeRef(json.withContents(f.field(["type"]).contents)), description: optString(f, "description"), deprecated: optString(f, "deprecated"), ) + let role = optString(f, "collection_role") + if (role == "keys") { field.withCollectionKeys(f.field(["name"]).asString) } + else if (role == "delta") { field.withCollectionDelta(f.field(["name"]).asString) } + else { field } } let withFunctions = raw.field(["functions"]).asArray.{{contents}}.reduce(withFields) { acc, rawFn => - acc.withFunction(functionOf(json.withContents(rawFn.contents))) + let f = json.withContents(rawFn.contents) + let withFunction = acc.withFunction(functionOf(f)) + if (f.field(["collection_get"]).asBoolean) { withFunction.withCollectionGet(f.field(["name"]).asString) } + else { withFunction } } let ctor = raw.field(["constructor"]) if (ctor.contents == "null") { diff --git a/sdk/src/dagger/client/gen.py b/sdk/src/dagger/client/gen.py index 8a852b2..a6d45a8 100644 --- a/sdk/src/dagger/client/gen.py +++ b/sdk/src/dagger/client/gen.py @@ -1625,6 +1625,78 @@ async def trace_url(self) -> str: return await _ctx.execute(str) +class CollectionDelta(Type): + async def added_keys(self) -> list[str]: + """Current keys absent from the original collection, in current order. + + Returns + ------- + list[str] + The `String` scalar type represents textual data, represented as + UTF-8 character sequences. The String type is most often used by + GraphQL to represent free-form human-readable text. + + Raises + ------ + ExecuteTimeoutError + If the time to execute the query exceeds the configured timeout. + QueryError + If the API returns an error. + """ + _args: list[Arg] = [] + _ctx = self._select("addedKeys", _args) + return await _ctx.execute(list[str]) + + async def id(self) -> str: + """A unique identifier for this CollectionDelta. + + Note + ---- + This is lazily evaluated, no operation is actually run. + + Returns + ------- + str + The `ID` scalar type represents a unique identifier, often used to + refetch an object or as key for a cache. The ID type appears in a + JSON response as a String; however, it is not intended to be + human-readable. When expected as an input type, any string (such + as `"4"`) or integer (such as `4`) input value will be accepted as + an ID. + + Raises + ------ + ExecuteTimeoutError + If the time to execute the query exceeds the configured timeout. + QueryError + If the API returns an error. + """ + _args: list[Arg] = [] + _ctx = self._select("id", _args) + return await _ctx.execute(str) + + async def removed_keys(self) -> list[str]: + """Original keys absent from the current collection, in original order. + + Returns + ------- + list[str] + The `String` scalar type represents textual data, represented as + UTF-8 character sequences. The String type is most often used by + GraphQL to represent free-form human-readable text. + + Raises + ------ + ExecuteTimeoutError + If the time to execute the query exceeds the configured timeout. + QueryError + If the API returns an error. + """ + _args: list[Arg] = [] + _ctx = self._select("removedKeys", _args) + return await _ctx.execute(list[str]) + + class Container(Type): """An OCI-compatible container, also known as a Docker container.""" @@ -16622,6 +16694,36 @@ async def optional(self) -> bool: _ctx = self._select("optional", _args) return await _ctx.execute(bool) + def with_collection(self) -> Self: + """Mark this object as a collection.""" + _args: list[Arg] = [] + _ctx = self._select("withCollection", _args) + return TypeDef(_ctx) + + def with_collection_delta(self, name: str) -> Self: + """Select the field that receives changes from the original collection.""" + _args = [ + Arg("name", name), + ] + _ctx = self._select("withCollectionDelta", _args) + return TypeDef(_ctx) + + def with_collection_get(self, name: str) -> Self: + """Select the item lookup function for this collection.""" + _args = [ + Arg("name", name), + ] + _ctx = self._select("withCollectionGet", _args) + return TypeDef(_ctx) + + def with_collection_keys(self, name: str) -> Self: + """Select the stored keys field for this collection.""" + _args = [ + Arg("name", name), + ] + _ctx = self._select("withCollectionKeys", _args) + return TypeDef(_ctx) + def with_constructor(self, function: Function) -> Self: """Adds a function for constructing a new instance of an Object TypeDef, failing if the type is not an object. @@ -19059,6 +19161,7 @@ class Client(Query): "Client", "ClientFilesyncMirror", "Cloud", + "CollectionDelta", "Container", "CurrentModule", "CurrentModuleAsSDK", diff --git a/sdk/src/dagger/mod/__init__.py b/sdk/src/dagger/mod/__init__.py index 1ab0293..73a5c92 100644 --- a/sdk/src/dagger/mod/__init__.py +++ b/sdk/src/dagger/mod/__init__.py @@ -13,6 +13,10 @@ agent = _default_mod.agent check = _default_mod.check +collection = _default_mod.collection +delta = _default_mod.delta +get = _default_mod.get +keys = _default_mod.keys enum_type = _default_mod.enum_type function = _default_mod.function field = _default_mod.field @@ -37,11 +41,15 @@ def default_module() -> Module: "Name", "agent", "check", + "collection", + "delta", "enum_type", "field", "function", "generate", + "get", "interface", + "keys", "object_type", "up", ] diff --git a/sdk/src/dagger/mod/_describe.py b/sdk/src/dagger/mod/_describe.py index ec5a488..8e5f85f 100644 --- a/sdk/src/dagger/mod/_describe.py +++ b/sdk/src/dagger/mod/_describe.py @@ -78,6 +78,7 @@ class FunctionDescription: generator: bool = False service: bool = False agent: bool = False + collection_get: bool = False args: tuple[ArgumentDescription, ...] = () @@ -88,6 +89,7 @@ class FieldDescription: description: str | None = None deprecated: str | None = None + collection_role: str | None = None @dataclasses.dataclass(frozen=True, slots=True) class EnumMemberDescription: @@ -108,6 +110,7 @@ class EnumDescription: class ObjectDescription: name: str interface: bool = False + collection: bool = False description: str | None = None deprecated: str | None = None fields: tuple[FieldDescription, ...] = () diff --git a/sdk/src/dagger/mod/_entrypoint.py b/sdk/src/dagger/mod/_entrypoint.py index 72035e4..5e8da31 100644 --- a/sdk/src/dagger/mod/_entrypoint.py +++ b/sdk/src/dagger/mod/_entrypoint.py @@ -123,13 +123,22 @@ def _object(obj: ObjectDescription) -> list[str]: deprecated = _opt("deprecated", obj.deprecated) head = f"typeDef.withObject({_quote(obj.name)}{description}{deprecated})" lines = [head] + if obj.collection: + lines.append(".withCollection") lines.extend( f".withField({_quote(field.name)}, {_type(field.type)}" f"{_opt('description', field.description)}" f"{_opt('deprecated', field.deprecated)})" for field in obj.fields ) + for field in obj.fields: + if field.collection_role == "keys": + lines.append(f".withCollectionKeys({_quote(field.name)})") + elif field.collection_role == "delta": + lines.append(f".withCollectionDelta({_quote(field.name)})") for func in obj.functions: + if func.collection_get: + lines.append(f".withCollectionGet({_quote(func.name)})") lines += _wrap("withFunction", _function(func)) if obj.constructor is not None: lines += _wrap("withConstructor", _function(obj.constructor)) diff --git a/sdk/src/dagger/mod/_module.py b/sdk/src/dagger/mod/_module.py index c0c515a..258f795 100644 --- a/sdk/src/dagger/mod/_module.py +++ b/sdk/src/dagger/mod/_module.py @@ -7,6 +7,7 @@ import textwrap import typing from collections.abc import Awaitable, Callable, Mapping +from functools import wraps from typing import Any, TypeVar, cast import anyio @@ -49,7 +50,14 @@ P, R, ) -from dagger.mod._types import APIName, FieldDefinition, FunctionDefinition, PythonName +from dagger.mod._types import ( + COLLECTION_BASE_ATTR, + COLLECTION_BASE_FIELD, + APIName, + FieldDefinition, + FunctionDefinition, + PythonName, +) from dagger.mod._utils import ( asyncify, extract_enum_member_doc, @@ -67,6 +75,8 @@ GENERATOR_DEF_KEY: typing.Final[str] = "__dagger_generate__" UP_DEF_KEY: typing.Final[str] = "__dagger_up__" AGENT_DEF_KEY: typing.Final[str] = "__dagger_agent__" +COLLECTION_DEF_KEY: typing.Final[str] = "__dagger_collection__" +COLLECTION_GET_DEF_KEY: typing.Final[str] = "__dagger_get__" MODULE_NAME: typing.Final[str] = os.getenv("DAGGER_MODULE", "") MAIN_OBJECT: typing.Final[str] = os.getenv("DAGGER_MAIN_OBJECT", "") TYPE_DEF_FILE: typing.Final[str] = os.getenv("DAGGER_MODULE_FILE", "/module.json") @@ -536,6 +546,51 @@ class Foo: **kwargs, ) + def collection(self, cls: T) -> T: + """Mark an exposed object type as a collection.""" + setattr(cls, COLLECTION_DEF_KEY, True) + return cls + + def get(self, func: Func[P, R]) -> Func[P, R]: + """Select the exposed item lookup function of a collection.""" + setattr(func, COLLECTION_GET_DEF_KEY, True) + return func + + def keys( + self, + *, + default: Callable[[], Any] | object = ..., + name: APIName | None = None, + init: bool = True, + deprecated: str | None = None, + ) -> Any: + """Expose the stored keys field of a collection.""" + return self._collection_field("keys", default, name, init, deprecated) + + def delta( + self, + *, + default: Callable[[], Any] | object = None, + name: APIName | None = None, + init: bool = False, + deprecated: str | None = None, + ) -> Any: + """Expose a field that receives the collection delta.""" + return self._collection_field("delta", default, name, init, deprecated) + + def _collection_field( + self, + role: str, + default: Callable[[], Any] | object, + name: APIName | None, + init: bool, + deprecated: str | None, + ) -> Any: + field = self.field(default=default, name=name, init=init, deprecated=deprecated) + meta = dataclasses.replace(field.metadata[FIELD_DEF_KEY], collection_role=role) + field.metadata = {**field.metadata, FIELD_DEF_KEY: meta} + return field + def check( self, func: Func[P, R] | None = None, @@ -736,14 +791,28 @@ def wrapper(func: Func[P, R]) -> Func[P, R]: @overload @dataclass_transform( kw_only_default=True, - field_specifiers=(function, dataclasses.field, dataclasses.Field), + field_specifiers=( + function, + field, + keys, + delta, + dataclasses.field, + dataclasses.Field, + ), ) def object_type(self, cls: T, /, *, deprecated: str | None = None) -> T: ... @overload @dataclass_transform( kw_only_default=True, - field_specifiers=(function, dataclasses.field, dataclasses.Field), + field_specifiers=( + function, + field, + keys, + delta, + dataclasses.field, + dataclasses.Field, + ), ) def object_type(self, *, deprecated: str | None = None) -> Callable[[T], T]: ... @@ -794,6 +863,25 @@ def wrapper(cls: T) -> T: ) raise BadUsageError(msg) + # Both decorator orders are supported: @collection can run after + # @object_type. Keep opaque engine state as an ordinary private field + # so copy.copy, deepcopy, and dataclasses.replace all preserve it. + cls.__annotations__ = dict(getattr(cls, "__annotations__", {})) + cls.__annotations__[COLLECTION_BASE_ATTR] = str | None + setattr( + cls, + COLLECTION_BASE_ATTR, + dataclasses.field(default=None, repr=False, compare=False), + ) + if init := cls.__dict__.get("__init__"): + + @wraps(init) + def init_with_state(instance, *args, **kwargs): + base = kwargs.pop(COLLECTION_BASE_ATTR, None) + init(instance, *args, **kwargs) + setattr(instance, COLLECTION_BASE_ATTR, base) + + cls.__init__ = init_with_state wrapped = dataclasses.dataclass(kw_only=True)(cls) return self._process_type(wrapped, deprecated=deprecated) @@ -838,7 +926,11 @@ def _is_function(fn) -> typing.TypeGuard[Func]: return cls # Register hooks for renaming field names in `mod.field()`. - attr_overrides = {} + attr_overrides = { + COLLECTION_BASE_ATTR: cattrs.gen.override( + rename=COLLECTION_BASE_FIELD, omit_if_default=True + ) + } # Find all fields exposed with `mod.field()`. for field in dataclasses.fields(cls): @@ -970,6 +1062,7 @@ def _describe_object(name: str, obj_type: ObjectType) -> ObjectDescription: ), description=get_doc(field.return_type), deprecated=field.meta.deprecated, + collection_role=field.meta.collection_role, ) for field in obj_type.fields.values() ) @@ -986,6 +1079,7 @@ def _describe_object(name: str, obj_type: ObjectType) -> ObjectDescription: return ObjectDescription( name=name, interface=obj_type.interface, + collection=getattr(obj_type.cls, COLLECTION_DEF_KEY, False), description=get_doc(obj_type.cls), deprecated=obj_type.deprecated, fields=fields, @@ -1028,6 +1122,7 @@ def _describe_function( generator=func.generate, service=func.service, agent=func.agent, + collection_get=getattr(func.wrapped, COLLECTION_GET_DEF_KEY, False), args=args, ) @@ -1075,6 +1170,8 @@ def _object_from(obj: ObjectDescription) -> dagger.TypeDef: description=obj.description, deprecated=obj.deprecated, ) + if obj.collection: + type_def = type_def.with_collection() for field in obj.fields: type_def = type_def.with_field( field.name, @@ -1082,7 +1179,13 @@ def _object_from(obj: ObjectDescription) -> dagger.TypeDef: description=field.description, deprecated=field.deprecated, ) + if field.collection_role == "keys": + type_def = type_def.with_collection_keys(field.name) + elif field.collection_role == "delta": + type_def = type_def.with_collection_delta(field.name) for func in obj.functions: + if func.collection_get: + type_def = type_def.with_collection_get(func.name) type_def = type_def.with_function(_function_from(func)) if obj.constructor is not None: type_def = type_def.with_constructor(_function_from(obj.constructor)) diff --git a/sdk/src/dagger/mod/_resolver.py b/sdk/src/dagger/mod/_resolver.py index f74f0ac..80b1275 100644 --- a/sdk/src/dagger/mod/_resolver.py +++ b/sdk/src/dagger/mod/_resolver.py @@ -22,7 +22,13 @@ InvalidInputError, RegistrationError, ) -from dagger.mod._types import APIName, FieldDefinition, FunctionDefinition, PythonName +from dagger.mod._types import ( + COLLECTION_BASE_ATTR, + APIName, + FieldDefinition, + FunctionDefinition, + PythonName, +) from dagger.mod._utils import ( get_alt_constructor, get_alt_name, @@ -155,7 +161,7 @@ def parameters(self): for param in self.signature.parameters.values(): # Skip `self` parameter on instance methods. # It will be added manually on `get_result`. - if param.name == "self": + if param.name in ("self", COLLECTION_BASE_ATTR): continue if param.kind is inspect.Parameter.POSITIONAL_ONLY: diff --git a/sdk/src/dagger/mod/_types.py b/sdk/src/dagger/mod/_types.py index 8e6e771..8860e29 100644 --- a/sdk/src/dagger/mod/_types.py +++ b/sdk/src/dagger/mod/_types.py @@ -8,12 +8,17 @@ APIName: TypeAlias = str ContextPath: TypeAlias = str +# Engine state. It is not part of the module schema. +COLLECTION_BASE_ATTR = "_dagger_collection_base" +COLLECTION_BASE_FIELD = "__daggerCollectionBase" + @dataclasses.dataclass(slots=True, frozen=True) class FieldDefinition: name: APIName | None optional: bool = False deprecated: str | None = None + collection_role: str | None = None @dataclasses.dataclass(slots=True, frozen=True) diff --git a/sdk/tests/mod/test_collections.py b/sdk/tests/mod/test_collections.py new file mode 100644 index 0000000..42256b0 --- /dev/null +++ b/sdk/tests/mod/test_collections.py @@ -0,0 +1,112 @@ +from copy import copy, deepcopy +from dataclasses import replace + +import pytest + +import dagger +from dagger.mod import Module + + +@pytest.mark.parametrize("collection_first", [True, False]) +@pytest.mark.parametrize("get_first", [True, False]) +def test_collection_markers_preserve_fields_and_functions(collection_first, get_first): + mod = Module() + + @mod.object_type + class Item: + name: str = mod.field() + + class Items: + names: list[str] = mod.keys(default=list, name="paths") + selection: dagger.CollectionDelta | None = mod.delta() + + def lookup(self, path: str) -> Item: + return Item(name=path) + + lookup = ( + mod.get(mod.function(lookup)) + if get_first + else mod.function(mod.get(lookup)) + ) + + items_type = ( + mod.collection(mod.object_type(Items)) + if collection_first + else mod.object_type(mod.collection(Items)) + ) + obj = mod.get_object("Items") + assert obj.cls.__dagger_collection__ + assert obj.fields["paths"].original_name == "names" + assert obj.fields["paths"].meta.collection_role == "keys" + assert obj.fields["selection"].meta.collection_role == "delta" + assert obj.functions["lookup"].wrapped.__dagger_get__ + assert items_type().names == [] + assert items_type().selection is None + + +@pytest.mark.parametrize("clone", [copy, deepcopy, replace]) +@pytest.mark.parametrize("custom_init", [False, True]) +def test_base_survives_copies_without_delta(clone, custom_init): + mod = Module() + + @mod.collection + @mod.object_type + class Items: + names: list[str] = mod.keys(default=list) + + if custom_init: + + def __init__(self, names: list[str]): + self.names = names + + original = mod._converter.structure( + {"names": ["a", "b"], "__daggerCollectionBase": "original"}, Items + ) + changed = clone(original) + changed.names = ["b", "c"] + assert mod._converter.unstructure(changed) == { + "names": ["b", "c"], + "__daggerCollectionBase": "original", + } + assert mod._converter.unstructure(Items(names=["c"])) == {"names": ["c"]} + assert "_dagger_collection_base" not in mod.get_object("Items").fields + assert ( + "_dagger_collection_base" + not in mod.get_object("Items").get_constructor().parameters + ) + + +def test_collection_metadata_reaches_both_entrypoints(): + from dagger.mod._entrypoint import render_types + from dagger.mod._module import _object_from + + mod = Module("Main") + + @mod.object_type + class Item: + name: str = mod.field() + + @mod.collection + @mod.object_type + class Main: + names: list[str] = mod.keys(default=list, name="paths") + selection: dagger.CollectionDelta | None = mod.delta() + + @mod.function(name="lookup") + @mod.get + def item(self, key: str) -> Item: + return Item(name=key) + + desc = mod.describe() + obj = next(obj for obj in desc.objects if obj.name == "Main") + assert obj.collection + assert next(f for f in obj.fields if f.name == "paths").collection_role == "keys" + assert next(f for f in obj.functions if f.name == "lookup").collection_get + code = render_types(desc) + assert ".withCollection" in code + assert '.withCollectionKeys("paths")' in code + assert '.withCollectionDelta("selection")' in code + assert '.withCollectionGet("lookup")' in code + # Build the runtime registration query without an engine connection. + runtime = _object_from(obj) + assert runtime is not None