bookshelf._produce#

class Activity:
    """Synchronous activity context that attributes every produced resource."""

    def __init__(
        self,
        client: BookshelfClient,
        cache: ContentCache,
        *,
        activity_id: UUID,
        kind: str,
        code_ref: str,
        config: Mapping[str, Any],
        runner: str,
        config_hash: str | None = None,
        default_visibility: models.Visibility = models.Visibility.hidden,
    ) -> None:
        self._client = client
        self._cache = cache
        self.activity_id = activity_id
        self.kind = kind
        self.code_ref = code_ref
        self.config = dict(config)
        self.runner = runner
        self.config_hash = config_hash
        self.default_visibility = default_visibility
        self._entered = False
        self._closed = False

    def __enter__(self) -> Self:
        if self._closed:
            raise RuntimeError("activity context cannot be re-entered after exit")
        if self._entered:
            raise RuntimeError("activity context is already entered")
        self._entered = True
        return self

    def __exit__(self, *exc_info: object) -> None:
        self._entered = False
        self._closed = True

    def _open(self) -> Self:
        """Make this activity usable outside a ``with`` block, for ``book.write``."""
        self._entered = True
        self._closed = False
        return self

    def _require_entered(self) -> None:
        if not self._entered:
            raise RuntimeError("register operations require an open 'with bs.activity(...)' block")

    def register(
        self,
        obj: object,
        *,
        type: str | models.ResourceType,
        name: str | None = None,
        used: Sequence[UsedInput] = (),
        visibility: VisibilityInput = INHERIT,
        tags: Sequence[str] = (),
        metadata: Mapping[str, Any] | None = None,
        tracking_id: UUID | None = None,
        format: str | None = None,
        dedupe: bool = True,
    ) -> Resource:
        """Serialise, hash, upload, and register one generated resource.

        ``used=`` records the inputs consumed by this resource.
        """
        self._require_entered()
        return self.register_many(
            [
                RegisterItem(
                    obj=obj,
                    type=type,
                    name=name,
                    visibility=visibility,
                    tags=tags,
                    metadata=metadata,
                    tracking_id=tracking_id,
                    format=format,
                    dedupe=dedupe,
                )
            ],
            used=used,
        )[0]

    def register_many(
        self,
        entries: Sequence[RegisterItem],
        *,
        used: Sequence[UsedInput] = (),
        atomic: bool = True,
    ) -> list[Resource]:
        """Materialise many outputs, splitting only a non atomic oversized batch.

        ``used=`` records the inputs consumed by every output in this call.
        """
        self._require_entered()
        if atomic and len(entries) > helpers.MAX_REGISTRATION_BATCH:
            raise ValueError(
                f"atomic registrations are limited to {helpers.MAX_REGISTRATION_BATCH} items"
            )
        items = [self._materialise(entry) for entry in entries]
        try:
            outcomes = self._register_items(items, used=used, atomic=atomic)
        except PartialRegistrationError as exc:
            exc.successful_resources = tuple(
                self._resource_from_success(success, items) for success in exc.successful
            )
            raise
        return [
            self._resource_from_outcome(outcome, item)
            for outcome, item in helpers.paired_successes(outcomes, items)
        ]

    def register_external(
        self,
        *,
        type: str | models.ResourceType,
        uri: str,
        hash: str | None = None,
        name: str | None = None,
        used: Sequence[UsedInput] = (),
        visibility: VisibilityInput = INHERIT,
        tags: Sequence[str] = (),
        metadata: Mapping[str, Any] | None = None,
        tracking_id: UUID | None = None,
        dedupe: bool = True,
    ) -> Resource:
        """Register an external output and attribute it to this activity.

        ``used=`` records the inputs consumed by this resource.
        """
        self._require_entered()
        item = helpers.external_item(
            type=type,
            uri=uri,
            hash=hash,
            name=name,
            visibility=helpers.visibility(visibility, self.default_visibility),
            tags=tags,
            metadata=metadata,
            tracking_id=tracking_id,
            dedupe=dedupe,
        )
        outcome = helpers.single_success(self._register_items([item], used=used, atomic=True))
        return Resource(
            self._client,
            self._cache,
            tracking_id=outcome.tracking_id,
            resource_type=helpers.registered_resource_type(outcome, item.type),
            registration_outcome=outcome,
            name=helpers.registered_name(item),
        )

    def _resource_from_outcome(
        self,
        outcome: models.RegistrationOutcome,
        item: models.RegisterResourceItem,
    ) -> Resource:
        return Resource(
            self._client,
            self._cache,
            tracking_id=outcome.tracking_id,
            resource_type=helpers.registered_resource_type(outcome, item.type),
            registration_outcome=outcome,
            name=helpers.registered_name(item),
        )

    def _resource_from_success(
        self,
        success: RegistrationSuccess,
        items: Sequence[models.RegisterResourceItem],
    ) -> Resource:
        return self._resource_from_outcome(success.outcome, items[success.index])

    def _materialise(self, entry: RegisterItem) -> models.RegisterResourceItem:
        resource_type = helpers.resource_type(entry.type)
        serialised = serialise(entry.obj, type=resource_type.value)
        storage_path = upload_bytes(
            self._client,
            serialised.data,
            hash_=serialised.hash,
            content_type=serialised.content_type,
        )
        return models.RegisterResourceItem(
            tracking_id=entry.tracking_id or helpers.uuid7(),
            type=resource_type,
            hash=serialised.hash,
            format=entry.format or serialised.format,
            name=entry.name,
            visibility=helpers.visibility(entry.visibility, self.default_visibility),
            discovery=helpers.resource_discovery(entry.tags),
            metadata=dict(entry.metadata or {}),
            locations=[models.LocationInput(shelf="managed", path=storage_path)],
            dedupe=entry.dedupe,
        )

    def _register_items(
        self,
        items: Sequence[models.RegisterResourceItem],
        *,
        used: Sequence[UsedInput],
        atomic: bool,
    ) -> list[RegistrationSuccess]:
        if atomic and len(items) > helpers.MAX_REGISTRATION_BATCH:
            raise ValueError(
                f"atomic registrations are limited to {helpers.MAX_REGISTRATION_BATCH} items"
            )
        # An atomic batch always goes in one request.
        chunk_size = max(len(items), 1) if atomic else helpers.MAX_REGISTRATION_BATCH
        successful: list[RegistrationSuccess] = []
        failures: list[RegistrationFailure] = []
        envelope = helpers.activity_envelope(
            activity_id=self.activity_id,
            kind=self.kind,
            code_ref=self.code_ref,
            config=self.config,
            runner=self.runner,
            used=used,
            config_hash=self.config_hash,
        )
        for start in range(0, len(items), chunk_size):
            response = self._client.register_resources(
                models.RegisterResourcesRequest(
                    items=list(items[start : start + chunk_size]),
                    activity=envelope,
                    atomic=atomic,
                )
            )
            chunk_successful, chunk_failures = helpers.registration_results(
                response,
                index_offset=start,
            )
            successful.extend(chunk_successful)
            failures.extend(chunk_failures)
        helpers.raise_partial_registration(successful, failures)
        return successful

class AsyncActivity:
    """Asynchronous activity context that attributes every produced resource."""

    def __init__(
        self,
        client: BookshelfClient,
        cache: ContentCache,
        *,
        activity_id: UUID,
        kind: str,
        code_ref: str,
        config: Mapping[str, Any],
        runner: str,
        config_hash: str | None = None,
        default_visibility: models.Visibility = models.Visibility.hidden,
    ) -> None:
        self._client = client
        self._cache = cache
        self.activity_id = activity_id
        self.kind = kind
        self.code_ref = code_ref
        self.config = dict(config)
        self.runner = runner
        self.config_hash = config_hash
        self.default_visibility = default_visibility
        self._entered = False
        self._closed = False

    async def __aenter__(self) -> Self:
        if self._closed:
            raise RuntimeError("activity context cannot be re-entered after exit")
        if self._entered:
            raise RuntimeError("activity context is already entered")
        self._entered = True
        return self

    async def __aexit__(self, *exc_info: object) -> None:
        self._entered = False
        self._closed = True

    def _open(self) -> Self:
        self._entered = True
        self._closed = False
        return self

    def _require_entered(self) -> None:
        if not self._entered:
            raise RuntimeError(
                "register operations require an open 'async with bs.activity(...)' block"
            )

    async def register(
        self,
        obj: object,
        *,
        type: str | models.ResourceType,
        name: str | None = None,
        used: Sequence[UsedInput] = (),
        visibility: VisibilityInput = INHERIT,
        tags: Sequence[str] = (),
        metadata: Mapping[str, Any] | None = None,
        tracking_id: UUID | None = None,
        format: str | None = None,
        dedupe: bool = True,
    ) -> AsyncResource:
        """Serialise, hash, upload, and register one generated resource.

        ``used=`` records the inputs consumed by this resource.
        """
        self._require_entered()
        resources = await self.register_many(
            [
                RegisterItem(
                    obj=obj,
                    type=type,
                    name=name,
                    visibility=visibility,
                    tags=tags,
                    metadata=metadata,
                    tracking_id=tracking_id,
                    format=format,
                    dedupe=dedupe,
                )
            ],
            used=used,
        )
        return resources[0]

    async def register_many(
        self,
        entries: Sequence[RegisterItem],
        *,
        used: Sequence[UsedInput] = (),
        atomic: bool = True,
    ) -> list[AsyncResource]:
        """Materialise many outputs, splitting only a non atomic oversized batch.

        ``used=`` records the inputs consumed by every output in this call.
        """
        self._require_entered()
        if atomic and len(entries) > helpers.MAX_REGISTRATION_BATCH:
            raise ValueError(
                f"atomic registrations are limited to {helpers.MAX_REGISTRATION_BATCH} items"
            )
        items: list[models.RegisterResourceItem] = []
        for entry in entries:
            items.append(await self._materialise(entry))
        try:
            outcomes = await self._register_items(items, used=used, atomic=atomic)
        except PartialRegistrationError as exc:
            exc.successful_resources = tuple(
                self._resource_from_success(success, items) for success in exc.successful
            )
            raise
        return [
            self._resource_from_outcome(outcome, item)
            for outcome, item in helpers.paired_successes(outcomes, items)
        ]

    async def register_external(
        self,
        *,
        type: str | models.ResourceType,
        uri: str,
        hash: str | None = None,
        name: str | None = None,
        used: Sequence[UsedInput] = (),
        visibility: VisibilityInput = INHERIT,
        tags: Sequence[str] = (),
        metadata: Mapping[str, Any] | None = None,
        tracking_id: UUID | None = None,
        dedupe: bool = True,
    ) -> AsyncResource:
        """Register an external output and attribute it to this activity.

        ``used=`` records the inputs consumed by this resource.
        """
        self._require_entered()
        item = helpers.external_item(
            type=type,
            uri=uri,
            hash=hash,
            name=name,
            visibility=helpers.visibility(visibility, self.default_visibility),
            tags=tags,
            metadata=metadata,
            tracking_id=tracking_id,
            dedupe=dedupe,
        )
        outcome = helpers.single_success(await self._register_items([item], used=used, atomic=True))
        return AsyncResource(
            self._client,
            self._cache,
            tracking_id=outcome.tracking_id,
            resource_type=helpers.registered_resource_type(outcome, item.type),
            registration_outcome=outcome,
            name=helpers.registered_name(item),
        )

    def _resource_from_outcome(
        self,
        outcome: models.RegistrationOutcome,
        item: models.RegisterResourceItem,
    ) -> AsyncResource:
        return AsyncResource(
            self._client,
            self._cache,
            tracking_id=outcome.tracking_id,
            resource_type=helpers.registered_resource_type(outcome, item.type),
            registration_outcome=outcome,
            name=helpers.registered_name(item),
        )

    def _resource_from_success(
        self,
        success: RegistrationSuccess,
        items: Sequence[models.RegisterResourceItem],
    ) -> AsyncResource:
        return self._resource_from_outcome(success.outcome, items[success.index])

    async def _materialise(self, entry: RegisterItem) -> models.RegisterResourceItem:
        resource_type = helpers.resource_type(entry.type)
        serialised = serialise(entry.obj, type=resource_type.value)
        storage_path = await upload_bytes_async(
            self._client,
            serialised.data,
            hash_=serialised.hash,
            content_type=serialised.content_type,
        )
        return models.RegisterResourceItem(
            tracking_id=entry.tracking_id or helpers.uuid7(),
            type=resource_type,
            hash=serialised.hash,
            format=entry.format or serialised.format,
            name=entry.name,
            visibility=helpers.visibility(entry.visibility, self.default_visibility),
            discovery=helpers.resource_discovery(entry.tags),
            metadata=dict(entry.metadata or {}),
            locations=[models.LocationInput(shelf="managed", path=storage_path)],
            dedupe=entry.dedupe,
        )

    async def _register_items(
        self,
        items: Sequence[models.RegisterResourceItem],
        *,
        used: Sequence[UsedInput],
        atomic: bool,
    ) -> list[RegistrationSuccess]:
        if atomic and len(items) > helpers.MAX_REGISTRATION_BATCH:
            raise ValueError(
                f"atomic registrations are limited to {helpers.MAX_REGISTRATION_BATCH} items"
            )
        # An atomic batch always goes in one request.
        chunk_size = max(len(items), 1) if atomic else helpers.MAX_REGISTRATION_BATCH
        successful: list[RegistrationSuccess] = []
        failures: list[RegistrationFailure] = []
        envelope = helpers.activity_envelope(
            activity_id=self.activity_id,
            kind=self.kind,
            code_ref=self.code_ref,
            config=self.config,
            runner=self.runner,
            used=used,
            config_hash=self.config_hash,
        )
        for start in range(0, len(items), chunk_size):
            response = await self._client.register_resources_async(
                models.RegisterResourcesRequest(
                    items=list(items[start : start + chunk_size]),
                    activity=envelope,
                    atomic=atomic,
                )
            )
            chunk_successful, chunk_failures = helpers.registration_results(
                response,
                index_offset=start,
            )
            successful.extend(chunk_successful)
            failures.extend(chunk_failures)
        helpers.raise_partial_registration(successful, failures)
        return successful

class AsyncDraftBook:
    """Mutable asynchronous draft-book handle."""

    def __init__(
        self,
        client: BookshelfClient,
        detail: models.BookDetail,
        *,
        activity: Callable[[], AsyncActivity] | None = None,
    ) -> None:
        self._client = client
        self.metadata = detail
        self._activity = activity

    def _writing_activity(self) -> AsyncActivity:
        if self._activity is None:
            raise RuntimeError(
                "book.write needs the activity its sink opens, and this book was drafted without one. "
                "Register through bs.activity(...) and attach with book.add."
            )
        return self._activity()

    async def write(
        self,
        name: str,
        obj: object,
        *,
        type: str | models.ResourceType = DEFAULT_WRITE_TYPE,
        used: Sequence[UsedInput] = (),
        data_dictionary: Sequence[models.DataDictionaryEntry] | None = None,
        visibility: VisibilityInput = INHERIT,
        tags: Sequence[str] = (),
        metadata: Mapping[str, Any] | None = None,
        format: str | None = None,
        dedupe: bool = True,
    ) -> Any:  # noqa: ANN401
        """Register one output and attach it under ``name`` in a single call.

        The asynchronous twin of :meth:`DraftBook.write`, with the same bundle result.
        """
        resource = await self._writing_activity().register(
            obj,
            type=type,
            name=name,
            used=used,
            visibility=visibility,
            tags=tags,
            metadata=metadata,
            format=format,
            dedupe=dedupe,
        )
        await self.attach(resource, name_in_book=name, data_dictionary=data_dictionary)
        return resource

    async def add(self, *resources: HasTrackingId) -> Self:
        """Attach already registered resources, each under the name it registered as."""
        for resource in resources:
            await self.attach(resource, name_in_book=_written_name(resource))
        return self

    @property
    def book_id(self) -> UUID:
        return self.metadata.book_id

    @property
    def status(self) -> str:
        return self.metadata.status

    async def attach(
        self,
        resource: HasTrackingId | str | UUID,
        *,
        name_in_book: str,
        data_dictionary: Sequence[models.DataDictionaryEntry] | None = None,
    ) -> models.BookEntryAttachResponse:
        """Attach a resource under a book-local name and optional entry dictionary.

        Omitting ``data_dictionary`` preserves the dictionary on an existing entry.
        Pass an empty sequence to clear it.
        """
        request = _attach_request(
            resource,
            name_in_book=name_in_book,
            data_dictionary=data_dictionary,
        )
        return await self._client.attach_entry_async(str(self.book_id), request)

    async def publish(self) -> Self:
        """Publish the assembled draft and update this handle in place."""
        self.metadata = await self._client.publish_book_async(str(self.book_id))
        return self