| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # pylint: disable=invalid-name |
| """Access path classes.""" |
| |
| # tvm-ffi-stubgen(begin): import-section |
| # fmt: off |
| # isort: off |
| from __future__ import annotations |
| from typing import TYPE_CHECKING |
| if TYPE_CHECKING: |
| from collections.abc import Sequence |
| from tvm_ffi import Object |
| from typing import Any |
| # isort: on |
| # fmt: on |
| # tvm-ffi-stubgen(end) |
| from enum import IntEnum |
| |
| from .core import Object |
| from .registry import register_object |
| |
| |
| class AccessKind(IntEnum): |
| """Kinds of access steps in an access path.""" |
| |
| ATTR = 0 |
| ARRAY_ITEM = 1 |
| MAP_ITEM = 2 |
| ATTR_MISSING = 3 |
| ARRAY_ITEM_MISSING = 4 |
| MAP_ITEM_MISSING = 5 |
| |
| |
| @register_object("ffi.reflection.AccessStep") |
| class AccessStep(Object): |
| """Access step container.""" |
| |
| # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep -> ffi.access_path.AccessStep |
| # tvm-ffi-stubgen(begin): object/ffi.reflection.AccessStep |
| # fmt: off |
| kind: int |
| key: Any |
| # fmt: on |
| # tvm-ffi-stubgen(end) |
| |
| |
| @register_object("ffi.reflection.AccessPath") |
| class AccessPath(Object): |
| """Access path container. |
| |
| It describes how to reach a nested attribute or item |
| inside a complex FFI object by recording a sequence of steps |
| (attribute, array index, or map key). It is primarily used by |
| diagnostics to pinpoint structural mismatches. |
| |
| Examples |
| -------- |
| .. code-block:: python |
| |
| from tvm_ffi.access_path import AccessPath |
| |
| root = AccessPath.root() |
| # Build a path equivalent to obj.layer.weight[2] |
| p = root.attr("layer").attr("weight").array_item(2) |
| assert isinstance(p, AccessPath) |
| |
| """ |
| |
| # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessPath -> ffi.access_path.AccessPath |
| # tvm-ffi-stubgen(begin): object/ffi.reflection.AccessPath |
| # fmt: off |
| parent: Object | None |
| step: AccessStep | None |
| depth: int |
| if TYPE_CHECKING: |
| @staticmethod |
| def _root() -> AccessPath: ... |
| def _extend(self, _1: AccessStep, /) -> AccessPath: ... |
| def _attr(self, _1: str, /) -> AccessPath: ... |
| def _array_item(self, _1: int, /) -> AccessPath: ... |
| def _map_item(self, _1: Any, /) -> AccessPath: ... |
| def _attr_missing(self, _1: str, /) -> AccessPath: ... |
| def _array_item_missing(self, _1: int, /) -> AccessPath: ... |
| def _map_item_missing(self, _1: Any, /) -> AccessPath: ... |
| def _is_prefix_of(self, _1: AccessPath, /) -> bool: ... |
| def _to_steps(self, /) -> Sequence[AccessStep]: ... |
| def _path_equal(self, _1: AccessPath, /) -> bool: ... |
| # fmt: on |
| # tvm-ffi-stubgen(end) |
| |
| def __init__(self) -> None: |
| """Disallow direct construction; use `AccessPath.root()` instead.""" |
| raise ValueError( |
| "AccessPath can't be initialized directly. " |
| "Use AccessPath.root() to create a path to the root object" |
| ) |
| |
| @staticmethod |
| def root() -> AccessPath: |
| """Create a root access path. |
| |
| Returns |
| ------- |
| AccessPath |
| A path representing the root of an object graph. |
| |
| """ |
| return AccessPath._root() |
| |
| def __eq__(self, other: Any) -> bool: |
| """Return whether two access paths are equal.""" |
| if not isinstance(other, AccessPath): |
| return False |
| return self._path_equal(other) |
| |
| def __ne__(self, other: Any) -> bool: |
| """Return whether two access paths are not equal.""" |
| if not isinstance(other, AccessPath): |
| return True |
| return not self._path_equal(other) |
| |
| def is_prefix_of(self, other: AccessPath) -> bool: |
| """Check if this access path is a prefix of another access path. |
| |
| Parameters |
| ---------- |
| other |
| The access path to check if it is a prefix of this access path |
| |
| Returns |
| ------- |
| bool |
| True if this access path is a prefix of the other access path, False otherwise |
| |
| """ |
| return self._is_prefix_of(other) |
| |
| def attr(self, attr_key: str) -> AccessPath: |
| """Create an access path to the attribute of the current object. |
| |
| Parameters |
| ---------- |
| attr_key |
| The key of the attribute to access |
| |
| Returns |
| ------- |
| AccessPath |
| The extended access path |
| |
| """ |
| return self._attr(attr_key) |
| |
| def attr_missing(self, attr_key: str) -> AccessPath: |
| """Create an access path that indicate an attribute is missing. |
| |
| Parameters |
| ---------- |
| attr_key |
| The key of the attribute to access |
| |
| Returns |
| ------- |
| AccessPath |
| The extended access path |
| |
| """ |
| return self._attr_missing(attr_key) |
| |
| def array_item(self, index: int) -> AccessPath: |
| """Create an access path to the item of the current array. |
| |
| Parameters |
| ---------- |
| index |
| The index of the item to access |
| |
| Returns |
| ------- |
| AccessPath |
| The extended access path |
| |
| """ |
| return self._array_item(index) |
| |
| def array_item_missing(self, index: int) -> AccessPath: |
| """Create an access path that indicate an array item is missing. |
| |
| Parameters |
| ---------- |
| index |
| The index of the item to access |
| |
| Returns |
| ------- |
| AccessPath |
| The extended access path |
| |
| """ |
| return self._array_item_missing(index) |
| |
| def map_item(self, key: Any) -> AccessPath: |
| """Create an access path to the item of the current map. |
| |
| Parameters |
| ---------- |
| key |
| The key of the item to access |
| |
| Returns |
| ------- |
| AccessPath |
| The extended access path |
| |
| """ |
| return self._map_item(key) |
| |
| def map_item_missing(self, key: Any) -> AccessPath: |
| """Create an access path that indicate a map item is missing. |
| |
| Parameters |
| ---------- |
| key |
| The key of the item to access |
| |
| Returns |
| ------- |
| AccessPath |
| The extended access path |
| |
| """ |
| return self._map_item_missing(key) |
| |
| def to_steps(self) -> Sequence[AccessStep]: |
| """Convert the access path to a list of access steps. |
| |
| Returns |
| ------- |
| access_steps |
| The list of access steps |
| |
| """ |
| return self._to_steps() |
| |
| __hash__ = Object.__hash__ |
