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

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"""Bridge-facing callback dispatch for Python ONNX Plugins.""" 

14 

15from ge.graph.operator import create_operator 

16 

17from .bootstrap import load_onnx_plugins 

18from ._native import OnnxNode 

19from .registry import ( 

20 get_registered_onnx_plugin_by_origin_type, 

21 get_registered_onnx_plugin_dicts, 

22) 

23 

24 

25class _InvalidParseNodeReturn(TypeError): 

26 """Internal marker for a parse_node callback returning a non-None value.""" 

27 

28 

29def load_and_get_onnx_plugin_descriptors() -> list: 

30 load_onnx_plugins() 

31 return get_registered_onnx_plugin_dicts() 

32 

33 

34def call_parse_node(origin_type: str, node: OnnxNode, operator_handle) -> None: 

35 """Dispatch one parser-owned ONNX node to its registered parse_node callback. 

36 

37 ``node`` and ``operator_handle`` are borrowed objects supplied by the C++ 

38 bridge for the duration of the callback. 

39 """ 

40 

41 descriptor = get_registered_onnx_plugin_by_origin_type(origin_type) 

42 if descriptor is None: 

43 raise KeyError(f"python ONNX Plugin is not registered: {origin_type}") 

44 

45 with create_operator(operator_handle) as target: 

46 result = descriptor.parser_node(node, target) 

47 if result is not None: 

48 raise _InvalidParseNodeReturn( 

49 "ONNX Plugin parse_node callback must return None" 

50 ) 

51 

52 

53def call_parse_operator(origin_type: str, source_handle, target_handle) -> None: 

54 """Dispatch one parser-owned Operator pair to its registered callback.""" 

55 

56 descriptor = get_registered_onnx_plugin_by_origin_type(origin_type) 

57 if descriptor is None: 

58 raise KeyError(f"python ONNX Plugin is not registered: {origin_type}") 

59 

60 with ( 

61 create_operator(source_handle, read_only=True) as source, 

62 create_operator(target_handle) as target, 

63 ): 

64 result = descriptor.parser_operator(source, target) 

65 if result is not None: 

66 raise _InvalidParseNodeReturn( 

67 "ONNX Plugin parse_operator callback must return None" 

68 )