Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/_signature.py: 89%
110 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-18 20:49 +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 annotations.pop("return", None)
111 if not annotations:
112 return {}
114 def annotation_source():
115 pass
117 annotation_source.__annotations__ = annotations
118 return typing.get_type_hints(
119 annotation_source,
120 globalns=getattr(target, "__globals__", None),
121 )
122 return typing.get_type_hints(method)
123 except (NameError, TypeError, AttributeError) as exc:
124 raise _signature_error(
125 descriptor,
126 method_name,
127 "resolvable type annotations",
128 f"type hint resolution failed: {exc}",
129 ) from exc
132def _validate_annotation(
133 parameter,
134 expected,
135 hints: dict,
136 descriptor,
137 method_name: str,
138 position: str,
139) -> None:
140 if parameter.annotation is inspect.Parameter.empty:
141 return
142 actual = hints.get(parameter.name, parameter.annotation)
143 if _normalize_annotation(actual) != _normalize_annotation(expected):
144 raise _signature_error(
145 descriptor,
146 method_name,
147 f"{position} annotation {_normalize_annotation(expected)!r}",
148 f"{_normalize_annotation(actual)!r}",
149 )
152def _validate_args_signature(
153 method,
154 ir_meta: dict,
155 descriptor,
156 *,
157 method_name: str = "declare_launch_args",
158) -> None:
159 if method_name not in ("execute", "declare_launch_args"):
160 raise ValueError(f"unsupported schema callback: {method_name}")
161 signature = inspect.signature(method)
162 parameters = list(signature.parameters.values())
163 for parameter in parameters:
164 if parameter.kind in (
165 inspect.Parameter.VAR_POSITIONAL,
166 inspect.Parameter.VAR_KEYWORD,
167 ):
168 raise _signature_error(
169 descriptor,
170 method_name,
171 "no variadic parameters",
172 f"variadic parameter {parameter.name}",
173 )
175 ir_inputs = ir_meta["inputs"]
176 ir_outputs = ir_meta["outputs"] if method_name == "declare_launch_args" else []
177 ir_attrs = ir_meta["attrs"]
178 positional_count = len(ir_inputs) + len(ir_outputs)
179 expected_count = positional_count + len(ir_attrs)
180 if len(parameters) != expected_count:
181 raise _signature_error(
182 descriptor,
183 method_name,
184 f"{positional_count} positional "
185 f"{'input/output' if method_name == 'declare_launch_args' else 'input'} "
186 "parameters followed by "
187 f"{len(ir_attrs)} keyword-only attrs",
188 f"{len(parameters)} parameters",
189 )
191 hints = _get_type_hints(method, descriptor, method_name)
192 for index, item in enumerate(ir_inputs):
193 parameter = parameters[index]
194 if parameter.kind not in _POSITIONAL_KINDS:
195 raise _signature_error(
196 descriptor,
197 method_name,
198 f"positional input parameter at index {index}",
199 f"parameter {parameter.name} kind {parameter.kind.name}",
200 )
201 _validate_annotation(
202 parameter,
203 _get_expected_input_annotation(item["kind"]),
204 hints,
205 descriptor,
206 method_name,
207 f"input parameter at index {index}",
208 )
210 for output_index, item in enumerate(ir_outputs):
211 parameter_index = len(ir_inputs) + output_index
212 parameter = parameters[parameter_index]
213 if parameter.kind not in _POSITIONAL_KINDS:
214 raise _signature_error(
215 descriptor,
216 method_name,
217 f"positional output parameter at index {output_index}",
218 f"parameter {parameter.name} kind {parameter.kind.name}",
219 )
220 _validate_annotation(
221 parameter,
222 _get_expected_output_annotation(item["kind"]),
223 hints,
224 descriptor,
225 method_name,
226 f"output parameter at index {output_index}",
227 )
229 for attr_index, item in enumerate(ir_attrs):
230 parameter = parameters[positional_count + attr_index]
231 if parameter.kind is not inspect.Parameter.KEYWORD_ONLY:
232 raise _signature_error(
233 descriptor,
234 method_name,
235 f"keyword-only attr parameter {item['name']}",
236 f"parameter {parameter.name} kind {parameter.kind.name}",
237 )
238 if parameter.name != item["name"]:
239 raise _signature_error(
240 descriptor,
241 method_name,
242 f"attr name {item['name']} at index {attr_index}",
243 f"attr name {parameter.name}",
244 )
245 _, expected_annotation = _get_runtime_attr_spec(item["type"], attr_index)
246 _validate_annotation(
247 parameter,
248 expected_annotation,
249 hints,
250 descriptor,
251 method_name,
252 f"attr parameter {item['name']}",
253 )
255 if method_name == "execute":
256 return
257 if signature.return_annotation is inspect.Signature.empty:
258 raise _signature_error(
259 descriptor,
260 method_name,
261 "None return annotation",
262 "missing return annotation",
263 )
264 return_annotation = hints.get("return", signature.return_annotation)
265 if _normalize_annotation(return_annotation) is not type(None):
266 raise _signature_error(
267 descriptor,
268 method_name,
269 "None return annotation",
270 repr(_normalize_annotation(return_annotation)),
271 )