-
Notifications
You must be signed in to change notification settings - Fork 4.5k
Expand file tree
/
Copy pathvalue_provider.py
More file actions
197 lines (159 loc) · 5.68 KB
/
value_provider.py
File metadata and controls
197 lines (159 loc) · 5.68 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""A ValueProvider abstracts the notion of fetching a value that may or
may not be currently available.
This can be used to parameterize transforms that only read values in at
runtime, for example.
"""
# pytype: skip-file
from functools import wraps
from typing import Set
from apache_beam import error
__all__ = [
'ValueProvider',
'StaticValueProvider',
'RuntimeValueProvider',
'NestedValueProvider',
'check_accessible',
]
class ValueProvider(object):
"""Base class that all other ValueProviders must implement.
"""
def is_accessible(self):
"""Whether the contents of this ValueProvider is available to routines
that run at graph construction time.
"""
raise NotImplementedError(
'ValueProvider.is_accessible implemented in derived classes')
def get(self):
"""Return the value wrapped by this ValueProvider.
"""
raise NotImplementedError(
'ValueProvider.get implemented in derived classes')
class StaticValueProvider(ValueProvider):
"""StaticValueProvider is an implementation of ValueProvider that allows
for a static value to be provided.
"""
def __init__(self, value_type, value):
"""
Args:
value_type: Type of the static value
value: Static value
"""
self.value_type = value_type
self.value = value_type(value)
def is_accessible(self):
return True
def get(self):
return self.value
def __str__(self):
return str(self.value)
def __eq__(self, other):
if self.value == other:
return True
if isinstance(other, StaticValueProvider):
if (self.value_type == other.value_type and self.value == other.value):
return True
return False
def __hash__(self):
return hash((type(self), self.value_type, self.value))
class RuntimeValueProvider(ValueProvider):
"""RuntimeValueProvider is an implementation of ValueProvider that
allows for a value to be provided at execution time rather than
at graph construction time.
"""
runtime_options = None
experiments: Set[str] = set()
def __init__(self, option_name, value_type, default_value):
self.option_name = option_name
self.default_value = default_value
self.value_type = value_type
def is_accessible(self):
return RuntimeValueProvider.runtime_options is not None
@classmethod
def get_value(cls, option_name, value_type, default_value):
if not RuntimeValueProvider.runtime_options:
return default_value
candidate = RuntimeValueProvider.runtime_options.get(option_name)
if candidate:
return value_type(candidate)
else:
return default_value
def get(self):
if RuntimeValueProvider.runtime_options is None:
raise error.RuntimeValueProviderError(
'%s.get() not called from a runtime context' % self)
return RuntimeValueProvider.get_value(
self.option_name, self.value_type, self.default_value)
@classmethod
def set_runtime_options(cls, pipeline_options):
RuntimeValueProvider.runtime_options = pipeline_options
RuntimeValueProvider.experiments = RuntimeValueProvider.get_value(
'experiments', set, set())
def __str__(self):
return '%s(option: %s, type: %s, default_value: %s)' % (
self.__class__.__name__,
self.option_name,
self.value_type.__name__,
repr(self.default_value))
class NestedValueProvider(ValueProvider):
"""NestedValueProvider is an implementation of ValueProvider that allows
for wrapping another ValueProvider object.
"""
def __init__(self, value, translator):
"""Creates a NestedValueProvider that wraps the provided ValueProvider.
Args:
value: ValueProvider object to wrap
translator: function that is applied to the ValueProvider
Raises:
``RuntimeValueProviderError``: if any of the provided objects are not
accessible.
"""
self.value = value
self.translator = translator
def is_accessible(self):
return self.value.is_accessible()
def get(self):
try:
return self.cached_value
except AttributeError:
self.cached_value = self.translator(self.value.get())
return self.cached_value
def __str__(self):
return "%s(value: %s, translator: %s)" % (
self.__class__.__name__,
self.value,
self.translator.__name__,
)
def check_accessible(value_provider_list):
"""A decorator that checks accessibility of a list of ValueProvider objects.
Args:
value_provider_list: list of ValueProvider objects
Raises:
``RuntimeValueProviderError``: if any of the provided objects are not
accessible.
"""
assert isinstance(value_provider_list, list)
def _check_accessible(fnc):
@wraps(fnc)
def _f(self, *args, **kwargs):
for obj in [getattr(self, vp) for vp in value_provider_list]:
if not obj.is_accessible():
raise error.RuntimeValueProviderError('%s not accessible' % obj)
return fnc(self, *args, **kwargs)
return _f
return _check_accessible