Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/dttest/api/python/ge/ge/onnx_plugin/registry.py: 98%

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"""Python ONNX Plugin descriptor registry.""" 

14 

15import threading 

16from dataclasses import dataclass, field 

17from typing import Callable, Dict, List, Optional, Tuple 

18 

19 

20@dataclass(frozen=True) 

21class OnnxPluginDescriptor: 

22 """Normalized descriptor for one parse_node callback.""" 

23 

24 descriptor_key: str 

25 source: str 

26 domain: str 

27 opsets: Tuple[int, ...] 

28 target: str 

29 origin_types: Tuple[str, ...] 

30 module_name: str 

31 parser_node: Callable[..., None] = field(compare=False, repr=False) 

32 

33 def to_bridge_dict(self) -> dict: 

34 return { 

35 "descriptor_key": self.descriptor_key, 

36 "source": self.source, 

37 "domain": self.domain, 

38 "opsets": list(self.opsets), 

39 "target": self.target, 

40 "origin_types": list(self.origin_types), 

41 "module_name": self.module_name, 

42 } 

43 

44 

45class _OnnxPluginRegistry: 

46 def __init__(self) -> None: 

47 self._lock = threading.RLock() 

48 self._descriptor_key_to_desc: Dict[str, OnnxPluginDescriptor] = {} 

49 self._origin_type_to_desc: Dict[str, OnnxPluginDescriptor] = {} 

50 

51 def clear(self) -> None: 

52 with self._lock: 

53 self._descriptor_key_to_desc.clear() 

54 self._origin_type_to_desc.clear() 

55 

56 def register(self, descriptor: OnnxPluginDescriptor) -> OnnxPluginDescriptor: 

57 with self._lock: 

58 if descriptor.descriptor_key in self._descriptor_key_to_desc: 

59 raise ValueError( 

60 "python ONNX Plugin descriptor_key already exists: " 

61 f"{descriptor.descriptor_key}" 

62 ) 

63 for origin_type in descriptor.origin_types: 

64 if origin_type in self._origin_type_to_desc: 

65 raise ValueError( 

66 f"python ONNX Plugin origin type already exists: {origin_type}" 

67 ) 

68 self._descriptor_key_to_desc[descriptor.descriptor_key] = descriptor 

69 for origin_type in descriptor.origin_types: 

70 self._origin_type_to_desc[origin_type] = descriptor 

71 return descriptor 

72 

73 def get_all(self) -> List[OnnxPluginDescriptor]: 

74 with self._lock: 

75 return list(self._descriptor_key_to_desc.values()) 

76 

77 def get_by_origin_type(self, origin_type: str) -> Optional[OnnxPluginDescriptor]: 

78 # Registration finishes before parsing; parse-time lookup is read-only. 

79 return self._origin_type_to_desc.get(origin_type) 

80 

81 

82_ONNX_PLUGIN_REGISTRY = _OnnxPluginRegistry() 

83 

84 

85def register_onnx_plugin( 

86 descriptor: OnnxPluginDescriptor, 

87) -> OnnxPluginDescriptor: 

88 return _ONNX_PLUGIN_REGISTRY.register(descriptor) 

89 

90 

91def clear_registered_onnx_plugins() -> None: 

92 _ONNX_PLUGIN_REGISTRY.clear() 

93 

94 

95def get_registered_onnx_plugins() -> List[OnnxPluginDescriptor]: 

96 return _ONNX_PLUGIN_REGISTRY.get_all() 

97 

98 

99def get_registered_onnx_plugin_dicts() -> List[dict]: 

100 return [item.to_bridge_dict() for item in get_registered_onnx_plugins()] 

101 

102 

103def get_registered_onnx_plugin_by_origin_type( 

104 origin_type: str, 

105) -> Optional[OnnxPluginDescriptor]: 

106 return _ONNX_PLUGIN_REGISTRY.get_by_origin_type(origin_type)