thds.core 0.0.1__py3-none-any.whl → 1.31.20250123022540__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of thds.core might be problematic. Click here for more details.

Files changed (70) hide show
  1. thds/core/__init__.py +48 -0
  2. thds/core/ansi_esc.py +46 -0
  3. thds/core/cache.py +201 -0
  4. thds/core/calgitver.py +82 -0
  5. thds/core/concurrency.py +100 -0
  6. thds/core/config.py +250 -0
  7. thds/core/decos.py +55 -0
  8. thds/core/dict_utils.py +188 -0
  9. thds/core/env.py +40 -0
  10. thds/core/exit_after.py +121 -0
  11. thds/core/files.py +125 -0
  12. thds/core/fretry.py +115 -0
  13. thds/core/generators.py +56 -0
  14. thds/core/git.py +81 -0
  15. thds/core/hash_cache.py +86 -0
  16. thds/core/hashing.py +106 -0
  17. thds/core/home.py +15 -0
  18. thds/core/hostname.py +10 -0
  19. thds/core/imports.py +17 -0
  20. thds/core/inspect.py +58 -0
  21. thds/core/iterators.py +9 -0
  22. thds/core/lazy.py +83 -0
  23. thds/core/link.py +153 -0
  24. thds/core/log/__init__.py +29 -0
  25. thds/core/log/basic_config.py +171 -0
  26. thds/core/log/json_formatter.py +43 -0
  27. thds/core/log/kw_formatter.py +84 -0
  28. thds/core/log/kw_logger.py +93 -0
  29. thds/core/log/logfmt.py +302 -0
  30. thds/core/merge_args.py +168 -0
  31. thds/core/meta.json +8 -0
  32. thds/core/meta.py +518 -0
  33. thds/core/parallel.py +200 -0
  34. thds/core/pickle_visit.py +24 -0
  35. thds/core/prof.py +276 -0
  36. thds/core/progress.py +112 -0
  37. thds/core/protocols.py +17 -0
  38. thds/core/py.typed +0 -0
  39. thds/core/scaling.py +39 -0
  40. thds/core/scope.py +199 -0
  41. thds/core/source.py +238 -0
  42. thds/core/source_serde.py +104 -0
  43. thds/core/sqlite/__init__.py +21 -0
  44. thds/core/sqlite/connect.py +33 -0
  45. thds/core/sqlite/copy.py +35 -0
  46. thds/core/sqlite/ddl.py +4 -0
  47. thds/core/sqlite/functions.py +63 -0
  48. thds/core/sqlite/index.py +22 -0
  49. thds/core/sqlite/insert_utils.py +23 -0
  50. thds/core/sqlite/merge.py +84 -0
  51. thds/core/sqlite/meta.py +190 -0
  52. thds/core/sqlite/read.py +66 -0
  53. thds/core/sqlite/sqlmap.py +179 -0
  54. thds/core/sqlite/structured.py +138 -0
  55. thds/core/sqlite/types.py +64 -0
  56. thds/core/sqlite/upsert.py +139 -0
  57. thds/core/sqlite/write.py +99 -0
  58. thds/core/stack_context.py +41 -0
  59. thds/core/thunks.py +40 -0
  60. thds/core/timer.py +214 -0
  61. thds/core/tmp.py +85 -0
  62. thds/core/types.py +4 -0
  63. thds.core-1.31.20250123022540.dist-info/METADATA +68 -0
  64. thds.core-1.31.20250123022540.dist-info/RECORD +67 -0
  65. {thds.core-0.0.1.dist-info → thds.core-1.31.20250123022540.dist-info}/WHEEL +1 -1
  66. thds.core-1.31.20250123022540.dist-info/entry_points.txt +4 -0
  67. thds.core-1.31.20250123022540.dist-info/top_level.txt +1 -0
  68. thds.core-0.0.1.dist-info/METADATA +0 -8
  69. thds.core-0.0.1.dist-info/RECORD +0 -4
  70. thds.core-0.0.1.dist-info/top_level.txt +0 -1
thds/core/config.py ADDED
@@ -0,0 +1,250 @@
1
+ """This is an attempt at a be-everything-to-everybody configuration 'system'.
2
+
3
+ Highlights:
4
+
5
+ - Configuration is always accessible and configurable via normal Python code.
6
+ - Configuration is type-safe.
7
+ - All active configuration is 'registered' and therefore discoverable.
8
+ - Config can be temporarily overridden for the current thread.
9
+ - Config can be set via a known environment variable.
10
+ - Config can be set by combining one or more configuration objects - these may be loaded from files,
11
+ but this system remains agnostic as to the format of those files or how and when they are actually loaded.
12
+
13
+ Please see thds/core/CONFIG.md for more details.
14
+ """
15
+
16
+ import typing as ty
17
+ from logging import getLogger
18
+ from os import getenv
19
+
20
+ from .stack_context import StackContext
21
+
22
+ _NOT_CONFIGURED = object()
23
+
24
+
25
+ class UnconfiguredError(ValueError):
26
+ pass
27
+
28
+
29
+ class ConfigNameCollisionError(KeyError):
30
+ pass
31
+
32
+
33
+ def _sanitize_env(env_var_name: str) -> str:
34
+ return env_var_name.replace("-", "_").replace(".", "_")
35
+
36
+
37
+ def _getenv(name: str, secret: bool) -> ty.Optional[str]:
38
+ """We want to support a variety of naming conventions for env
39
+ vars, without requiring people to actually name their config using
40
+ all caps and underscores only.
41
+
42
+ Many modern shells support more complex env var names.
43
+ """
44
+
45
+ def _first_nonnil(*envvars: str) -> ty.Optional[str]:
46
+ for envvar in envvars:
47
+ raw_value = getenv(envvar)
48
+ if raw_value is not None:
49
+ lvalue = "***SECRET***" if secret else raw_value
50
+ getLogger(__name__).info(
51
+ f"Loaded config '{name}' with raw value '{lvalue}' from environment variable '{envvar}'"
52
+ )
53
+ return raw_value
54
+ return None
55
+
56
+ return _first_nonnil(name, _sanitize_env(name), _sanitize_env(name).upper())
57
+
58
+
59
+ def _module_name():
60
+ import inspect
61
+
62
+ # this takes ~200 nanoseconds to run. that's fast enough that
63
+ # doing the extra lookup is worth it for being more user-friendly.
64
+ frame = inspect.currentframe()
65
+ frame = frame.f_back.f_back.f_back # type: ignore
66
+ return frame.f_globals["__name__"] # type: ignore
67
+
68
+
69
+ def _fullname(name: str) -> str:
70
+ """In the vast majority of cases, the best way to use this library
71
+ is to name your configuration items to correspond with your
72
+ fully-qualified module name as a prefix.
73
+
74
+ It will enhance discoverability and clarity,
75
+ and will help avoid configuration name collisions.
76
+ """
77
+ return name if "." in name else f"{_module_name()}.{name}"
78
+
79
+
80
+ T = ty.TypeVar("T")
81
+
82
+
83
+ def _type_parser(default: T) -> ty.Callable[[ty.Any], T]:
84
+ if default is None:
85
+ return lambda x: x # cannot learn anything about how to parse a future value from None
86
+ try:
87
+ default_type = type(default)
88
+ if default == default_type(default): # type: ignore
89
+ # if this succeeds and doesn't raise, then the type is self-parsing.
90
+ # in other words, type(4)(4) == 4, type('foobar')('foobar') == 'foobar'
91
+ return default_type
92
+ except Exception:
93
+ pass
94
+ return lambda x: x # we can't infer a type parser, so we'll return the default no-op parser.
95
+
96
+
97
+ class ConfigItem(ty.Generic[T]):
98
+ """Should only ever be constructed at a module level."""
99
+
100
+ def __init__(
101
+ self,
102
+ name: str,
103
+ # names must be hierarchical. if name contains no `.`, we will force
104
+ # your name to be in the calling module.
105
+ default: T = ty.cast(T, _NOT_CONFIGURED),
106
+ *,
107
+ parse: ty.Optional[ty.Callable[[ty.Any], T]] = None,
108
+ secret: bool = False,
109
+ ):
110
+ self.secret = secret
111
+ name = _fullname(name)
112
+ if name in _REGISTRY:
113
+ raise ConfigNameCollisionError(f"Config item {name} has already been registered!")
114
+ _REGISTRY[name] = self
115
+ self.name = name
116
+ self.parse = parse or _type_parser(default)
117
+ raw_resolved_global = _getenv(name, secret=secret)
118
+ if raw_resolved_global:
119
+ # external global values are only resolved at initial
120
+ # creation. if you want to set this value globally after
121
+ # application start, use set_global.
122
+ self.global_value = self.parse(raw_resolved_global)
123
+ elif default is not _NOT_CONFIGURED:
124
+ self.global_value = self.parse(default)
125
+ else:
126
+ self.global_value = default
127
+ self._stack_context: StackContext[T] = StackContext(
128
+ "config_" + name, ty.cast(T, _NOT_CONFIGURED)
129
+ )
130
+
131
+ def set_global(self, value: T):
132
+ """Global to the current process.
133
+
134
+ Will not automatically get transferred to spawned processes.
135
+ """
136
+
137
+ self.global_value = value
138
+ # we used to parse this value, but I think that was the wrong choice -
139
+ # we should only have to parse it when it's coming from the environment,
140
+ # which should be handled by it's initial creation,
141
+ # or when it's being set as a global default from a config file.
142
+
143
+ def set_local(self, value: T) -> ty.ContextManager[T]:
144
+ """Local to the current thread.
145
+
146
+ Will not automatically get transferred to spawned threads.
147
+ """
148
+
149
+ return self._stack_context.set(value)
150
+
151
+ def __call__(self) -> T:
152
+ local = self._stack_context()
153
+ if local is not _NOT_CONFIGURED:
154
+ return local
155
+ if self.global_value is _NOT_CONFIGURED:
156
+ raise UnconfiguredError(f"Config item '{self.name}' has not been configured!")
157
+ return self.global_value
158
+
159
+
160
+ def tobool(s_or_b: ty.Union[str, bool]) -> bool:
161
+ """A reasonable implementation that we could expand in the future."""
162
+ return s_or_b if isinstance(s_or_b, bool) else s_or_b.lower() not in ("0", "false", "no", "off", "")
163
+
164
+
165
+ def maybe(parser: ty.Callable[[ty.Any], T]) -> ty.Callable[[ty.Optional[ty.Any]], ty.Optional[T]]:
166
+ """A helper for when you want to parse a value that might be nil."""
167
+ return lambda x: parser(x) if x is not None else None
168
+
169
+
170
+ item = ConfigItem
171
+ # a short alias
172
+
173
+
174
+ _REGISTRY: ty.Dict[str, ConfigItem] = dict()
175
+
176
+
177
+ def config_by_name(name: str) -> ConfigItem:
178
+ """This is a dynamic interface - in general, prefer accessing the ConfigItem object directly."""
179
+ return _REGISTRY[_fullname(name)]
180
+
181
+
182
+ def set_global_defaults(config: ty.Dict[str, ty.Any]):
183
+ """Any config-file parser can create a dictionary of only the
184
+ items it managed to read, and then all of those can be set at once
185
+ via this function.
186
+ """
187
+ for name, value in config.items():
188
+ if not isinstance(value, dict):
189
+ try:
190
+ config_item = _REGISTRY[name]
191
+ config_item.set_global(config_item.parse(value))
192
+ except KeyError:
193
+ # try directly importing a module - this is only best-effort and will not work
194
+ # if you did not follow standard configuration naming conventions.
195
+ import importlib
196
+
197
+ maybe_module_name = ".".join(name.split(".")[:-1])
198
+ try:
199
+ importlib.import_module(maybe_module_name)
200
+ try:
201
+ config_item = _REGISTRY[name]
202
+ config_item.set_global(config_item.parse(value))
203
+ except KeyError as kerr:
204
+ raise KeyError(
205
+ f"Config item {name} is not registered"
206
+ f" and no module with the name {maybe_module_name} was importable."
207
+ " Please double-check your configuration."
208
+ ) from kerr
209
+ except ModuleNotFoundError:
210
+ # create a new, dynamic config item that will only be accessible via
211
+ # its name.
212
+ ConfigItem(name, value) # return value not needed since it self-registers.
213
+ getLogger(__name__).debug(
214
+ "Created dynamic config item '%s' with value '%s'", name, value
215
+ )
216
+
217
+ else: # recurse
218
+ set_global_defaults({f"{name}.{key}": val for key, val in value.items()})
219
+
220
+
221
+ def get_all_config() -> ty.Dict[str, ty.Any]:
222
+ return {k: v() if not v.secret else "***SECRET***" for k, v in _REGISTRY.items()}
223
+
224
+
225
+ def show_config_cli():
226
+ import argparse
227
+ import importlib
228
+ from pprint import pprint
229
+
230
+ parser = argparse.ArgumentParser()
231
+ parser.add_argument("via_modules", type=str, nargs="+")
232
+ args = parser.parse_args()
233
+
234
+ for module in args.via_modules:
235
+ importlib.import_module(module)
236
+
237
+ print()
238
+ print("thds.core.config")
239
+ print("----------------")
240
+ print("The following keys are fully-qualified module paths by default.")
241
+ print(
242
+ "A given item can be capitalized and set via environment variable"
243
+ ", e.g. export THDS_CORE_LOG_LEVEL='DEBUG'"
244
+ )
245
+ print(
246
+ "An item can also be set globally or locally, after importing the"
247
+ " ConfigItem object from the Python module where it is defined."
248
+ )
249
+ print()
250
+ pprint(get_all_config())
thds/core/decos.py ADDED
@@ -0,0 +1,55 @@
1
+ import typing as ty
2
+
3
+ F = ty.TypeVar("F", bound=ty.Callable)
4
+
5
+
6
+ @ty.overload
7
+ def compose(*decorators: ty.Callable[[F], F]) -> ty.Callable[[F], F]:
8
+ ... # pragma: no cover
9
+
10
+
11
+ @ty.overload
12
+ def compose(*decorators: ty.Callable[[F], F], f: F) -> F:
13
+ ... # pragma: no cover
14
+
15
+
16
+ def compose(*decorators: ty.Callable[[F], F], f: ty.Optional[F] = None) -> ty.Callable[[F], F]:
17
+ """A decorator factory that creates a single decorator from
18
+ multiple, following the top-to-bottom order that would be used by
19
+ the `@deco` decorator syntax, which is actually standard R-to-L
20
+ composition order.
21
+
22
+ It may also be used to both compose the decorator and apply it
23
+ directly to a provided function.
24
+
25
+ This is useful when you want to apply a sequence of decorators
26
+ (higher-order functions) to a function, but you can't use the
27
+ decorator syntax because you want the original function to be
28
+ callable without the decorators.
29
+
30
+ Example:
31
+
32
+ >>> decos.compose_partial(
33
+ ... deco1,
34
+ ... deco2,
35
+ ... deco3,
36
+ ... )(f)
37
+
38
+ is equivalent to:
39
+
40
+ >>> @deco1
41
+ ... @deco2
42
+ ... @deco3
43
+ ... def f():
44
+ ... pass
45
+
46
+ """
47
+
48
+ def _deco(func: F) -> F:
49
+ for deco in reversed(decorators):
50
+ func = deco(func)
51
+ return func
52
+
53
+ if f:
54
+ return _deco(f)
55
+ return _deco
@@ -0,0 +1,188 @@
1
+ import re
2
+ import warnings
3
+ from collections import defaultdict
4
+ from typing import Any, Dict, Generator, List, Mapping, MutableMapping, Optional, Tuple, TypeVar
5
+
6
+ DEFAULT_SEP = "."
7
+ VT = TypeVar("VT")
8
+
9
+
10
+ def _get_valid_variable_name(var: str):
11
+ """
12
+ given a string returns the string formatted as a proper python variable name.
13
+ Credit: https://stackoverflow.com/questions/3303312/how-do-i-convert-a-string-to-a-valid-variable-name-in-python
14
+ """
15
+ return re.sub(r"\W+|^(?=\d)", "_", var)
16
+
17
+
18
+ def _flatten_gen(
19
+ d: Mapping, parent_key: str = "", sep: str = DEFAULT_SEP
20
+ ) -> Generator[Tuple[str, Any], None, None]:
21
+ """
22
+ flattens a mapping (usually a dict) using a separator, returning a generator of the flattened keys and values.
23
+
24
+ Example
25
+ ---------
26
+
27
+ d = {"a": {"b": {"c": 1}}}
28
+ fd = flatten(d, sep=".")
29
+ print(dict(fd))
30
+ > {"a.b.c": 1}
31
+ """
32
+ for k, v in d.items():
33
+ new_key = parent_key + sep + k if parent_key else k
34
+ if isinstance(v, Mapping):
35
+ yield from _flatten_gen(v, new_key, sep=sep)
36
+ else:
37
+ yield new_key, v
38
+
39
+
40
+ def unflatten(flat_d: Dict[str, Any], sep: str = DEFAULT_SEP):
41
+ """Given a flattened dictionary returns the un-flatten representation."""
42
+ unflatten_dict: Dict[str, Any] = {}
43
+ for path, val in flat_d.items():
44
+ dict_ref = unflatten_dict
45
+ path_parts = path.split(sep)
46
+ for p in path_parts[:-1]:
47
+ dict_ref[p] = dict_ref.get(p) or {}
48
+ dict_ref = dict_ref[p]
49
+ dict_ref[path_parts[-1]] = val
50
+ return unflatten_dict
51
+
52
+
53
+ def flatten(d: Mapping, parent_key: str = "", sep: str = DEFAULT_SEP) -> Dict[str, Any]:
54
+ return dict(_flatten_gen(d, parent_key, sep))
55
+
56
+
57
+ class DotDict(MutableMapping[str, VT]):
58
+ """A python dictionary that acts like an object."""
59
+
60
+ _new_to_orig_keys: Dict[str, str] = dict()
61
+ _hidden_data: Dict[str, Any] = dict()
62
+
63
+ def _get_hidden_data(self, identifier: str) -> Any:
64
+ return self._hidden_data.get(identifier)
65
+
66
+ def _construct(self, mapping: Mapping) -> None:
67
+ convert_keys_to_identifiers = self._get_hidden_data("convert_keys_to_identifiers")
68
+ for k, v in mapping.items():
69
+ new_key = _get_valid_variable_name(k) if convert_keys_to_identifiers else k
70
+ if convert_keys_to_identifiers:
71
+ self._new_to_orig_keys[new_key] = k
72
+ if isinstance(v, dict):
73
+ self[new_key] = DotDict(v) # type: ignore
74
+ elif isinstance(v, (list, tuple, set)):
75
+ self[new_key] = v.__class__([DotDict(iv) if isinstance(iv, dict) else iv for iv in v]) # type: ignore
76
+ else:
77
+ self[new_key] = v
78
+
79
+ def __init__(self, *args, convert_keys_to_identifiers: bool = False, **kwargs):
80
+ self._hidden_data["convert_keys_to_identifiers"] = convert_keys_to_identifiers
81
+ if convert_keys_to_identifiers:
82
+ warnings.warn("automatically converting keys into identifiers. Data loss might occur.")
83
+ for arg in args:
84
+ if isinstance(arg, dict):
85
+ self._construct(mapping=arg)
86
+ if kwargs:
87
+ self._construct(mapping=kwargs)
88
+
89
+ def __getattr__(self, key: str) -> VT:
90
+ return self[key]
91
+
92
+ def __setattr__(self, key: str, value: VT) -> None:
93
+ self.__setitem__(key, value)
94
+
95
+ def __setitem__(self, key: str, value: VT):
96
+ self.__dict__.update({key: value})
97
+
98
+ def __delattr__(self, key: str) -> None:
99
+ self.__delitem__(key)
100
+
101
+ def __delitem__(self, key: str) -> None:
102
+ super(DotDict, self).__delitem__(key)
103
+ del self.__dict__[key]
104
+
105
+ def __getitem__(self, key: str) -> VT:
106
+ return self.__dict__[key]
107
+
108
+ def __iter__(self):
109
+ return iter(self.__dict__)
110
+
111
+ def __len__(self) -> int:
112
+ return len(self.__dict__)
113
+
114
+ def to_dict(self, orig_keys: bool = False) -> Dict[str, VT]:
115
+ convert_keys_to_identifiers = self._get_hidden_data("convert_keys_to_identifiers")
116
+ d: Dict[str, VT] = dict()
117
+ for k, v in self.items():
118
+ if isinstance(v, DotDict):
119
+ d[
120
+ self._new_to_orig_keys[k] if orig_keys and convert_keys_to_identifiers else k
121
+ ] = v.to_dict(
122
+ orig_keys
123
+ ) # type: ignore[assignment]
124
+ else:
125
+ d[self._new_to_orig_keys[k] if orig_keys and convert_keys_to_identifiers else k] = v
126
+ return d
127
+
128
+ def get_value(self, dot_path: str) -> Optional[VT]:
129
+ """Get a value given a dotted path to the value.
130
+
131
+ Example
132
+ -------
133
+
134
+ dd = DotDict(a={"b": 100})
135
+ assert dd.get_value("a.b") == 100
136
+ """
137
+ path = dot_path.split(".")
138
+ ref: DotDict[Any] = self
139
+ for k in path[:-1]:
140
+ if isinstance(ref, DotDict) and k in ref:
141
+ ref = ref[k]
142
+ else:
143
+ return None
144
+ try:
145
+ return ref[path[-1]]
146
+ except KeyError:
147
+ return None
148
+
149
+ def set_value(self, dot_path: str, val: VT) -> None:
150
+ """Set a vlaue given a dotted path."""
151
+ ref = self
152
+ path = dot_path.split(".")
153
+ try:
154
+ for k in path[:-1]:
155
+ ref = getattr(ref, k)
156
+ except AttributeError:
157
+ raise KeyError(f"can't set path {dot_path} with parts {path}.")
158
+ ref.__setattr__(path[-1], val)
159
+
160
+
161
+ def merge_dicts(*dicts: Dict[Any, Any], default: Any = None) -> Dict[Any, Any]:
162
+ """Merges similar dictionaries into one dictionary where the resulting values are a list of values from the
163
+ original dicts. If a dictionary does not have a key the default value will be used (defaults to None).
164
+
165
+ Example
166
+ --------
167
+
168
+ assert merge_dicts(
169
+ {"a": 100, "b": {"c": 200, "d": 300}, "e": [1, 2]},
170
+ {"a": 200, "b": {"c": 300}, "e": [3, 4], "f": 300}
171
+ ) == {
172
+ "a": [100, 200],
173
+ "b": {
174
+ "c": [200, 300],
175
+ "d": [300, None]
176
+ },
177
+ "e": [[1,2], [3,4]],
178
+ "f": [None, 300]
179
+ }
180
+ """
181
+ merged_dict: Dict[str, List[Any]] = defaultdict(lambda: [default for _ in range(len(dicts))])
182
+ for i, d in enumerate(dicts):
183
+ for k, v in d.items() if isinstance(d, dict) else {}:
184
+ if isinstance(v, dict):
185
+ merged_dict[k] = merge_dicts(*[a.get(k, {}) for a in dicts]) # type: ignore
186
+ else:
187
+ merged_dict[k][i] = v
188
+ return dict(merged_dict)
thds/core/env.py ADDED
@@ -0,0 +1,40 @@
1
+ import os
2
+ import typing as ty
3
+ from contextlib import contextmanager
4
+
5
+ from .stack_context import StackContext
6
+
7
+ THDS_ENV = "THDS_ENV"
8
+
9
+ Env = ty.Literal["", "prod", "dev", "ua-prod"]
10
+
11
+ _TEMP_ENV: StackContext[Env] = StackContext("thds-env", "")
12
+
13
+
14
+ def set_active_env(env: Env):
15
+ os.environ[THDS_ENV] = env
16
+
17
+
18
+ def get_raw_env() -> Env:
19
+ """Get the actual value of `THDS_ENV`. Unset == ''
20
+
21
+ Prefer `active_env` for determining the active environment.
22
+ """
23
+ return ty.cast(Env, os.environ.get(THDS_ENV, ""))
24
+
25
+
26
+ @contextmanager
27
+ def temp_env(env: Env = "") -> ty.Iterator[Env]:
28
+ """Temporarily set the value of `THDS_ENV` for the current stack/thread.
29
+
30
+ Useful if you have special cases where you need to fetch a result
31
+ from a different environment than the one you're intending to run,
32
+ without affecting the global state of your application.
33
+ """
34
+ with _TEMP_ENV.set(env or active_env()):
35
+ yield _TEMP_ENV()
36
+
37
+
38
+ def active_env(override: Env = "") -> Env:
39
+ """Get the effective value of `THDS_ENV`."""
40
+ return override or _TEMP_ENV() or get_raw_env() or "dev"
@@ -0,0 +1,121 @@
1
+ """A program that exits after a given number of seconds.
2
+
3
+ but whose exit can be delayed by modifying a file on the filesystem.
4
+ """
5
+
6
+ import argparse
7
+ import os
8
+ import time
9
+ import typing as ty
10
+ from datetime import datetime, timedelta, timezone
11
+ from logging import getLogger
12
+ from pathlib import Path
13
+
14
+ logger = getLogger(__name__)
15
+
16
+
17
+ def now() -> datetime:
18
+ return datetime.now(tz=timezone.utc)
19
+
20
+
21
+ def report(exit_time: datetime) -> datetime:
22
+ reporting_at = now()
23
+ msg = f"Will exit at {exit_time}, in {exit_time - reporting_at}"
24
+ logger.info(msg)
25
+ return reporting_at
26
+
27
+
28
+ def try_parse(time_file: Path, default: datetime) -> datetime:
29
+ dt_s = ""
30
+ try:
31
+ with open(time_file) as f:
32
+ dt_s = f.read().strip()
33
+ dt = datetime.fromisoformat(dt_s)
34
+ if default != dt:
35
+ logger.info(f"Parsed new time {dt} from {time_file}")
36
+ return dt
37
+ except FileNotFoundError:
38
+ logger.debug(f"No file found at {time_file}")
39
+ except ValueError:
40
+ logger.exception(f"Unable to parse {time_file} with contents {dt_s}; using default {default}")
41
+ return default
42
+
43
+
44
+ _1_SECOND = timedelta(seconds=1)
45
+
46
+
47
+ def exit_when(
48
+ exit_time_file: Path,
49
+ exit_at: datetime,
50
+ *,
51
+ report_every: ty.Optional[timedelta],
52
+ check_every: timedelta = _1_SECOND,
53
+ keep_existing_file: bool = False,
54
+ ) -> timedelta:
55
+ # setup
56
+ started = now()
57
+ exit_time_file = exit_time_file.resolve()
58
+ if not exit_at.tzinfo:
59
+ exit_at = exit_at.astimezone(timezone.utc)
60
+
61
+ if not keep_existing_file:
62
+ with open(exit_time_file, "w") as f:
63
+ f.write(exit_at.isoformat())
64
+
65
+ exit_at = try_parse(exit_time_file, exit_at)
66
+ if report_every:
67
+ reported_at = report(exit_at)
68
+ while True:
69
+ rem = (exit_at - now()).total_seconds()
70
+ if rem <= 0:
71
+ break
72
+ time.sleep(min(check_every.total_seconds(), rem))
73
+ exit_at = try_parse(exit_time_file, exit_at)
74
+ if report_every and now() - reported_at > report_every:
75
+ reported_at = report(exit_at)
76
+
77
+ os.unlink(exit_time_file)
78
+ return now() - started
79
+
80
+
81
+ def main():
82
+ parser = argparse.ArgumentParser(description=__doc__)
83
+ parser.add_argument(
84
+ "--exit-time-file",
85
+ "-f",
86
+ default="~/WILL_EXIT_AT_THIS_TIME.txt",
87
+ help=(
88
+ "Modify this file while the program is running to accelerate or delay the exit."
89
+ " Will be deleted and recreated if it exists, unless --keep-existing-file is also passed."
90
+ ),
91
+ )
92
+ parser.add_argument(
93
+ "--seconds-from-now",
94
+ "-s",
95
+ default=timedelta(hours=1).total_seconds(),
96
+ type=float,
97
+ )
98
+ parser.add_argument(
99
+ "--report-every-s",
100
+ "-r",
101
+ default=timedelta(minutes=20).total_seconds(),
102
+ type=float,
103
+ help="Report time of exit this often (in seconds)",
104
+ )
105
+ parser.add_argument(
106
+ "--keep-existing-file",
107
+ action="store_true",
108
+ help="Unless set, an existing file will be presumed to be from a previous run",
109
+ )
110
+ args = parser.parse_args()
111
+ elapsed = exit_when(
112
+ Path(os.path.expanduser(args.exit_time_file)),
113
+ datetime.now(tz=timezone.utc) + timedelta(seconds=args.seconds_from_now),
114
+ report_every=timedelta(seconds=args.report_every_s),
115
+ keep_existing_file=args.keep_existing_file,
116
+ )
117
+ logger.info(f"Exiting after {elapsed}")
118
+
119
+
120
+ if __name__ == "__main__":
121
+ main()