Set up thread-local state and return a contextmanager for managing it. Args: name: string, converted to lowercase to define the name of the config option (and absl flag). It is converted to uppercase to define the corresponding shell environment variable.
(
self,
name: str,
enum_values: List[str],
default: Optional[str],
help: str,
update_global_hook: Optional[Callable[[str], None]] = None,
update_thread_local_hook: Optional[Callable[[Optional[str]], None]] = None,
)
| 224 | ) |
| 225 | |
| 226 | def define_enum_state( |
| 227 | self, |
| 228 | name: str, |
| 229 | enum_values: List[str], |
| 230 | default: Optional[str], |
| 231 | help: str, |
| 232 | update_global_hook: Optional[Callable[[str], None]] = None, |
| 233 | update_thread_local_hook: Optional[Callable[[Optional[str]], None]] = None, |
| 234 | ): |
| 235 | """Set up thread-local state and return a contextmanager for managing it. |
| 236 | Args: |
| 237 | name: string, converted to lowercase to define the name of the config |
| 238 | option (and absl flag). It is converted to uppercase to define the |
| 239 | corresponding shell environment variable. |
| 240 | enum_values: list of strings representing the possible values for the |
| 241 | option. |
| 242 | default: optional string, default value. |
| 243 | help: string, used to populate the flag help information as well as the |
| 244 | docstring of the returned context manager. |
| 245 | Returns: |
| 246 | A contextmanager to control the thread-local state value. |
| 247 | See docstring for ``define_bool_state``. |
| 248 | """ |
| 249 | name = name.lower() |
| 250 | default = os.getenv(name.upper(), default) |
| 251 | if default is not None and default not in enum_values: |
| 252 | raise ValueError(f'Invalid value "{default}" for XLA flag {name}') |
| 253 | self.DEFINE_enum( |
| 254 | name, |
| 255 | default, |
| 256 | enum_values=enum_values, |
| 257 | help=help, |
| 258 | update_hook=update_global_hook, |
| 259 | ) |
| 260 | self._contextmanager_flags.add(name) |
| 261 | |
| 262 | def get_state(self): |
| 263 | val = _thread_local_state.__dict__.get(name, unset) |
| 264 | return val if val is not unset else self._read(name) |
| 265 | |
| 266 | setattr(Config, name, property(get_state)) |
| 267 | |
| 268 | def validate(new_val): |
| 269 | if new_val is not None and ( |
| 270 | type(new_val) is not str or new_val not in enum_values |
| 271 | ): |
| 272 | raise ValueError( |
| 273 | f"new enum value must be None or in {enum_values}, " |
| 274 | f"got {new_val} of type {type(new_val)}." |
| 275 | ) |
| 276 | |
| 277 | return _StateContextManager(name, help, update_thread_local_hook, validate) |
| 278 | |
| 279 | def define_int_state( |
| 280 | self, |
no test coverage detected