Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/custom_op/base.py: 97%

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# ----------------------------------------------------------------------------------------------------------- 

12 

13"""Base definitions for Python GE custom ops.""" 

14 

15from __future__ import annotations 

16import inspect 

17from abc import ABC, abstractmethod 

18from functools import wraps 

19from typing import Callable, List 

20 

21from ._native import EagerOpExecutionContext 

22 

23 

24_MISSING = object() 

25_POSITIONAL_PARAMETER_KINDS = ( 

26 inspect.Parameter.POSITIONAL_ONLY, 

27 inspect.Parameter.POSITIONAL_OR_KEYWORD, 

28) 

29_EXECUTE_CONTEXT_ARG_NAME = "ctx" 

30_EXECUTE_INPUTS_ARG_NAME = "inputs" 

31_SUPPORTED_EXECUTE_ARG_NAMES = (_EXECUTE_CONTEXT_ARG_NAME, _EXECUTE_INPUTS_ARG_NAME) 

32 

33 

34def _build_input_tensor_list(ctx: EagerOpExecutionContext) -> List[object]: 

35 return [ctx.get_input_tensor(index) for index in range(ctx.get_input_num())] 

36 

37 

38def _inject_ctx_methods(instance, ctx: EagerOpExecutionContext) -> dict: 

39 original_attrs = {} 

40 for name in dir(ctx): 

41 if name.startswith("_"): 

42 continue 

43 ctx_method = getattr(ctx, name) 

44 if not callable(ctx_method): 

45 continue 

46 original_attrs[name] = instance.__dict__.get(name, _MISSING) 

47 setattr(instance, name, ctx_method) 

48 return original_attrs 

49 

50 

51def _restore_ctx_methods(instance, original_attrs: dict) -> None: 

52 for name, value in original_attrs.items(): 

53 if value is _MISSING: 

54 instance.__dict__.pop(name, None) 

55 else: 

56 setattr(instance, name, value) 

57 

58 

59def _adapt_inputs_execute(method: Callable) -> Callable: 

60 @wraps(method) 

61 def wrapper(self, ctx: EagerOpExecutionContext) -> None: 

62 inputs = _build_input_tensor_list(ctx) 

63 original_attrs = _inject_ctx_methods(self, ctx) 

64 try: 

65 method(self, inputs) 

66 finally: 

67 _restore_ctx_methods(self, original_attrs) 

68 

69 return wrapper 

70 

71 

72def _build_execute_signature_error(method: Callable) -> str: 

73 return ( 

74 "EagerExecuteOp.execute only supports execute(self, ctx) or " 

75 f"execute(self, inputs); got execute{inspect.signature(method)}. " 

76 "Use 'ctx' for EagerOpExecutionContext or 'inputs' for the input tensor list." 

77 ) 

78 

79 

80def _positional_params(method: Callable) -> List[inspect.Parameter]: 

81 return [ 

82 param 

83 for param in inspect.signature(method).parameters.values() 

84 if param.kind in _POSITIONAL_PARAMETER_KINDS 

85 ] 

86 

87 

88def _has_non_positional_params(method: Callable) -> bool: 

89 return any( 

90 param.kind not in _POSITIONAL_PARAMETER_KINDS 

91 for param in inspect.signature(method).parameters.values() 

92 ) 

93 

94 

95def _check_execute_arg_count(method: Callable) -> None: 

96 if len(_positional_params(method)) != 2 or _has_non_positional_params(method): 

97 raise TypeError(_build_execute_signature_error(method)) 

98 

99 

100def _get_execute_arg_name(method: Callable) -> str: 

101 _check_execute_arg_count(method) 

102 arg_name = _positional_params(method)[1].name 

103 if arg_name not in _SUPPORTED_EXECUTE_ARG_NAMES: 

104 raise TypeError(_build_execute_signature_error(method)) 

105 return arg_name 

106 

107 

108class BaseCustomOp(ABC): 

109 """Base class for Python custom ops.""" 

110 

111 

112class EagerExecuteOp(BaseCustomOp): 

113 """Base class for Python eager execute custom ops.""" 

114 

115 def __init_subclass__(cls, **kwargs): 

116 super().__init_subclass__(**kwargs) 

117 method = cls.__dict__.get("execute") 

118 if callable(method): 

119 arg_name = _get_execute_arg_name(method) 

120 if arg_name == _EXECUTE_INPUTS_ARG_NAME: 

121 cls.execute = _adapt_inputs_execute(method) 

122 

123 @abstractmethod 

124 def execute(self, ctx: EagerOpExecutionContext) -> None: 

125 raise NotImplementedError