Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/_signature.py: 90%
107 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 10:20 +0800
1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3# -----------------------------------------------------------------------------------------------------------
4# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6# CANN Open Software License Agreement Version 2.0 (the "License").
7# Please refer to the License for details. You may not use this file except in compliance with the License.
8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10# See LICENSE in the root of the software repository for the full text of the License.
11# -----------------------------------------------------------------------------------------------------------
13"""Schema-bound callback signature validation and runtime attribute metadata."""
15import inspect
16import types
17import typing
19from ge.graph import DataType
20from ge.runtime import Tensor
22from ._ir_types import AttrType, InputType, OutputType
25_POSITIONAL_KINDS = (
26 inspect.Parameter.POSITIONAL_ONLY,
27 inspect.Parameter.POSITIONAL_OR_KEYWORD,
28)
29_GET_ORIGIN = getattr(
30 typing, "get_origin", lambda value: getattr(value, "__origin__", None)
31)
32_GET_ARGS = getattr(typing, "get_args", lambda value: getattr(value, "__args__", ()))
33_UNION_ORIGINS = {typing.Union}
34_PEP604_UNION = getattr(types, "UnionType", None)
35if _PEP604_UNION is not None:
36 _UNION_ORIGINS.add(_PEP604_UNION)
38_RUNTIME_ATTR_SPECS = {
39 AttrType.INT: ("get_int", int),
40 AttrType.FLOAT: ("get_float", float),
41 AttrType.BOOL: ("get_bool", bool),
42 AttrType.STRING: ("get_str", str),
43 AttrType.DATA_TYPE: ("get_data_type", DataType),
44 AttrType.TENSOR: ("get_tensor", Tensor),
45 AttrType.LIST_INT: ("get_list_int", list[int]),
46 AttrType.LIST_FLOAT: ("get_list_float", list[float]),
47 AttrType.LIST_BOOL: ("get_list_bool", list[bool]),
48 AttrType.LIST_STRING: ("get_list_str", list[str]),
49 AttrType.LIST_DATA_TYPE: ("get_list_data_type", list[DataType]),
50 AttrType.LIST_LIST_INT: ("get_list_list_int", list[list[int]]),
51}
54def _signature_error(
55 descriptor, method_name: str, expected: str, actual: str
56) -> TypeError:
57 return TypeError(
58 f"invalid {method_name} signature for op type "
59 f"{descriptor.op_type}, descriptor key {descriptor.descriptor_key}, "
60 f"method {method_name}: expected {expected}, actual {actual}"
61 )
64def _normalize_annotation(annotation):
65 if annotation is None:
66 return type(None)
67 origin = _GET_ORIGIN(annotation)
68 args = _GET_ARGS(annotation)
69 if origin is list:
70 return ("list", tuple(_normalize_annotation(arg) for arg in args))
71 if origin in _UNION_ORIGINS:
72 return ("union", frozenset(_normalize_annotation(arg) for arg in args))
73 return annotation
76def _get_expected_input_annotation(kind: int):
77 if kind == InputType.REQUIRED:
78 return Tensor
79 if kind == InputType.OPTIONAL:
80 return typing.Optional[Tensor]
81 if kind == InputType.DYNAMIC:
82 return list[Tensor]
83 raise ValueError(f"unsupported custom op IR input kind: {kind}")
86def _get_expected_output_annotation(kind: int):
87 if kind == OutputType.REQUIRED:
88 return Tensor
89 if kind == OutputType.DYNAMIC:
90 return list[Tensor]
91 raise ValueError(f"unsupported custom op IR output kind: {kind}")
94def _get_runtime_attr_spec(ir_type: str, index: int):
95 spec = _RUNTIME_ATTR_SPECS.get(ir_type)
96 if spec is None:
97 raise ValueError(
98 f"unsupported custom op runtime attr type: {ir_type}, attr index: {index}"
99 )
100 return spec
103def _get_type_hints(method, descriptor, method_name: str) -> dict:
104 try:
105 if getattr(method, "__no_type_check__", False):
106 return {}
107 if method_name == "execute":
108 target = getattr(method, "__func__", method)
109 annotations = dict(getattr(target, "__annotations__", {}))
110 if not annotations:
111 return {}
113 def annotation_source():
114 pass
116 annotation_source.__annotations__ = annotations
117 return typing.get_type_hints(
118 annotation_source,
119 globalns=getattr(target, "__globals__", None),
120 )
121 return typing.get_type_hints(method)
122 except (NameError, TypeError, AttributeError) as exc:
123 raise _signature_error(
124 descriptor,
125 method_name,
126 "resolvable type annotations",
127 f"type hint resolution failed: {exc}",
128 ) from exc
131def _validate_annotation(
132 parameter,
133 expected,
134 hints: dict,
135 descriptor,
136 method_name: str,
137 position: str,
138) -> None:
139 if parameter.annotation is inspect.Parameter.empty:
140 return
141 actual = hints.get(parameter.name, parameter.annotation)
142 if _normalize_annotation(actual) != _normalize_annotation(expected):
143 raise _signature_error(
144 descriptor,
145 method_name,
146 f"{position} annotation {_normalize_annotation(expected)!r}",
147 f"{_normalize_annotation(actual)!r}",
148 )
151def _validate_args_signature(
152 method,
153 ir_meta: dict,
154 descriptor,
155 *,
156 method_name: str = "declare_launch_args",
157) -> None:
158 if method_name not in ("execute", "declare_launch_args"):
159 raise ValueError(f"unsupported schema callback: {method_name}")
160 signature = inspect.signature(method)
161 parameters = list(signature.parameters.values())
162 for parameter in parameters:
163 if parameter.kind in (
164 inspect.Parameter.VAR_POSITIONAL,
165 inspect.Parameter.VAR_KEYWORD,
166 ):
167 raise _signature_error(
168 descriptor,
169 method_name,
170 "no variadic parameters",
171 f"variadic parameter {parameter.name}",
172 )
174 ir_inputs = ir_meta["inputs"]
175 ir_outputs = ir_meta["outputs"] if method_name == "declare_launch_args" else []
176 ir_attrs = ir_meta["attrs"]
177 positional_count = len(ir_inputs) + len(ir_outputs)
178 expected_count = positional_count + len(ir_attrs)
179 if len(parameters) != expected_count:
180 raise _signature_error(
181 descriptor,
182 method_name,
183 f"{positional_count} positional "
184 f"{'input/output' if method_name == 'declare_launch_args' else 'input'} "
185 "parameters followed by "
186 f"{len(ir_attrs)} keyword-only attrs",
187 f"{len(parameters)} parameters",
188 )
190 hints = _get_type_hints(method, descriptor, method_name)
191 for index, item in enumerate(ir_inputs):
192 parameter = parameters[index]
193 if parameter.kind not in _POSITIONAL_KINDS:
194 raise _signature_error(
195 descriptor,
196 method_name,
197 f"positional input parameter at index {index}",
198 f"parameter {parameter.name} kind {parameter.kind.name}",
199 )
200 _validate_annotation(
201 parameter,
202 _get_expected_input_annotation(item["kind"]),
203 hints,
204 descriptor,
205 method_name,
206 f"input parameter at index {index}",
207 )
209 for output_index, item in enumerate(ir_outputs):
210 parameter_index = len(ir_inputs) + output_index
211 parameter = parameters[parameter_index]
212 if parameter.kind not in _POSITIONAL_KINDS:
213 raise _signature_error(
214 descriptor,
215 method_name,
216 f"positional output parameter at index {output_index}",
217 f"parameter {parameter.name} kind {parameter.kind.name}",
218 )
219 _validate_annotation(
220 parameter,
221 _get_expected_output_annotation(item["kind"]),
222 hints,
223 descriptor,
224 method_name,
225 f"output parameter at index {output_index}",
226 )
228 for attr_index, item in enumerate(ir_attrs):
229 parameter = parameters[positional_count + attr_index]
230 if parameter.kind is not inspect.Parameter.KEYWORD_ONLY:
231 raise _signature_error(
232 descriptor,
233 method_name,
234 f"keyword-only attr parameter {item['name']}",
235 f"parameter {parameter.name} kind {parameter.kind.name}",
236 )
237 if parameter.name != item["name"]:
238 raise _signature_error(
239 descriptor,
240 method_name,
241 f"attr name {item['name']} at index {attr_index}",
242 f"attr name {parameter.name}",
243 )
244 _, expected_annotation = _get_runtime_attr_spec(item["type"], attr_index)
245 _validate_annotation(
246 parameter,
247 expected_annotation,
248 hints,
249 descriptor,
250 method_name,
251 f"attr parameter {item['name']}",
252 )
254 if signature.return_annotation is inspect.Signature.empty:
255 raise _signature_error(
256 descriptor,
257 method_name,
258 "None return annotation",
259 "missing return annotation",
260 )
261 return_annotation = hints.get("return", signature.return_annotation)
262 if _normalize_annotation(return_annotation) is not type(None):
263 raise _signature_error(
264 descriptor,
265 method_name,
266 "None return annotation",
267 repr(_normalize_annotation(return_annotation)),
268 )