Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/keras/src/saving/legacy/saved_model/load_context.py: 61%

1# Copyright 2020 The TensorFlow Authors. All Rights Reserved. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================== 

15"""Context for storing options for loading a SavedModel.""" 

16 

17import contextlib 

18import threading 

19 

20import tensorflow.compat.v2 as tf 

21 

22 

23class LoadContext(threading.local): 

24 """A context for loading a model.""" 

25 

26 def __init__(self): 

27 super().__init__() 

28 self._entered_load_context = [] 

29 self._load_options = None 

30 

31 def set_load_options(self, load_options): 

32 self._load_options = load_options 

33 self._entered_load_context.append(True) 

34 

35 def clear_load_options(self): 

36 self._load_options = None 

37 self._entered_load_context.pop() 

38 

39 def load_options(self): 

40 return self._load_options 

41 

42 def in_load_context(self): 

43 return self._entered_load_context 

44 

45 

46_load_context = LoadContext() 

47 

48 

49@contextlib.contextmanager 

50def load_context(load_options): 

51 _load_context.set_load_options(load_options) 

52 try: 

53 yield 

54 finally: 

55 _load_context.clear_load_options() 

56 

57 

58def get_load_options(): 

59 """Returns the load options under a load context.""" 

60 return _load_context.load_options() 

61 

62 

63def in_load_context(): 

64 """Returns whether under a load context.""" 

65 return _load_context.in_load_context() 

66 

67 

68tf.__internal__.register_load_context_function(in_load_context) 

69