Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load_context.py: 59%

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 

20 

21class LoadContext(threading.local): 

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

23 

24 def __init__(self): 

25 super(LoadContext, self).__init__() 

26 self._entered_load_context = [] 

27 self._load_options = None 

28 

29 def set_load_options(self, load_options): 

30 self._load_options = load_options 

31 self._entered_load_context.append(True) 

32 

33 def clear_load_options(self): 

34 self._load_options = None 

35 self._entered_load_context.pop() 

36 

37 def load_options(self): 

38 return self._load_options 

39 

40 def in_load_context(self): 

41 return self._entered_load_context 

42 

43 

44_load_context = LoadContext() 

45 

46 

47@contextlib.contextmanager 

48def load_context(load_options): 

49 _load_context.set_load_options(load_options) 

50 try: 

51 yield 

52 finally: 

53 _load_context.clear_load_options() 

54 

55 

56def get_load_options(): 

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

58 return _load_context.load_options() 

59 

60 

61def in_load_context(): 

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

63 return _load_context.in_load_context()