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

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"""Execution context access for schema-bound Python custom ops.""" 

14 

15from contextlib import contextmanager 

16from contextvars import ContextVar 

17from dataclasses import dataclass 

18from typing import Iterator, Optional 

19 

20from ._native import EagerOpExecutionContext 

21 

22 

23@dataclass 

24class _ExecuteContextBinding: 

25 ctx: EagerOpExecutionContext 

26 active: bool = True 

27 

28 

29_CURRENT_EXECUTE_CONTEXT: ContextVar[Optional[_ExecuteContextBinding]] = ContextVar( 

30 "ge_custom_op_execute_context", default=None 

31) 

32 

33 

34def get_execute_ctx() -> EagerOpExecutionContext: 

35 """Return the borrowed context of the active schema-bound execute callback.""" 

36 

37 binding = _CURRENT_EXECUTE_CONTEXT.get() 

38 if binding is None or not binding.active: 

39 raise RuntimeError( 

40 "get_execute_ctx() is only available inside schema-bound execute" 

41 ) 

42 return binding.ctx 

43 

44 

45@contextmanager 

46def _execute_ctx_scope(ctx: EagerOpExecutionContext) -> Iterator[None]: 

47 binding = _ExecuteContextBinding(ctx=ctx) 

48 token = _CURRENT_EXECUTE_CONTEXT.set(binding) 

49 try: 

50 yield 

51 finally: 

52 binding.active = False 

53 _CURRENT_EXECUTE_CONTEXT.reset(token)