ominfra 0.0.0.dev134__py3-none-any.whl → 0.0.0.dev135__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
@@ -0,0 +1,1635 @@
1
+ #!/usr/bin/env python3
2
+ # noinspection DuplicatedCode
3
+ # @omlish-lite
4
+ # @omlish-script
5
+ # @omlish-amalg-output main.py
6
+ # ruff: noqa: N802 UP006 UP007 UP036
7
+ import abc
8
+ import base64
9
+ import collections.abc
10
+ import contextlib
11
+ import dataclasses as dc
12
+ import datetime
13
+ import decimal
14
+ import enum
15
+ import fractions
16
+ import functools
17
+ import inspect
18
+ import json
19
+ import logging
20
+ import os
21
+ import platform
22
+ import pwd
23
+ import shlex
24
+ import site
25
+ import struct
26
+ import subprocess
27
+ import sys
28
+ import textwrap
29
+ import threading
30
+ import time
31
+ import types
32
+ import typing as ta
33
+ import uuid
34
+ import weakref # noqa
35
+ import zlib
36
+
37
+
38
+ ########################################
39
+
40
+
41
+ if sys.version_info < (3, 8):
42
+ raise OSError(f'Requires python (3, 8), got {sys.version_info} from {sys.executable}') # noqa
43
+
44
+
45
+ ########################################
46
+
47
+
48
+ # commands/base.py
49
+ CommandInputT = ta.TypeVar('CommandInputT', bound='Command.Input')
50
+ CommandOutputT = ta.TypeVar('CommandOutputT', bound='Command.Output')
51
+
52
+ # ../../../omlish/lite/cached.py
53
+ T = ta.TypeVar('T')
54
+
55
+ # ../../../omlish/lite/check.py
56
+ SizedT = ta.TypeVar('SizedT', bound=ta.Sized)
57
+
58
+
59
+ ########################################
60
+ # ../commands/base.py
61
+
62
+
63
+ ##
64
+
65
+
66
+ class Command(abc.ABC, ta.Generic[CommandInputT, CommandOutputT]):
67
+ @dc.dataclass(frozen=True)
68
+ class Input(abc.ABC): # noqa
69
+ pass
70
+
71
+ @dc.dataclass(frozen=True)
72
+ class Output(abc.ABC): # noqa
73
+ pass
74
+
75
+ @abc.abstractmethod
76
+ def _execute(self, inp: CommandInputT) -> CommandOutputT:
77
+ raise NotImplementedError
78
+
79
+
80
+ ########################################
81
+ # ../../../pyremote.py
82
+ """
83
+ Basically this: https://mitogen.networkgenomics.com/howitworks.html
84
+ """
85
+
86
+
87
+ ##
88
+
89
+
90
+ _PYREMOTE_BOOTSTRAP_COMM_FD = 100
91
+ _PYREMOTE_BOOTSTRAP_SRC_FD = 101
92
+
93
+ _PYREMOTE_BOOTSTRAP_CHILD_PID_VAR = '_OPYR_CPID'
94
+ _PYREMOTE_BOOTSTRAP_ARGV0_VAR = '_OPYR_ARGV0'
95
+
96
+ _PYREMOTE_BOOTSTRAP_ACK0 = b'OPYR000\n'
97
+ _PYREMOTE_BOOTSTRAP_ACK1 = b'OPYR001\n'
98
+ _PYREMOTE_BOOTSTRAP_ACK2 = b'OPYR002\n'
99
+ _PYREMOTE_BOOTSTRAP_ACK3 = b'OPYR003\n'
100
+
101
+ _PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT = '(pyremote:%s)'
102
+
103
+ _PYREMOTE_BOOTSTRAP_IMPORTS = [
104
+ 'base64',
105
+ 'os',
106
+ 'struct',
107
+ 'sys',
108
+ 'zlib',
109
+ ]
110
+
111
+
112
+ def _pyremote_bootstrap_main(context_name: str) -> None:
113
+ # Get pid
114
+ pid = os.getpid()
115
+
116
+ # Two copies of main src to be sent to parent
117
+ r0, w0 = os.pipe()
118
+ r1, w1 = os.pipe()
119
+
120
+ if (cp := os.fork()):
121
+ # Parent process
122
+
123
+ # Dup original stdin to comm_fd for use as comm channel
124
+ os.dup2(0, _PYREMOTE_BOOTSTRAP_COMM_FD)
125
+
126
+ # Overwrite stdin (fed to python repl) with first copy of src
127
+ os.dup2(r0, 0)
128
+
129
+ # Dup second copy of src to src_fd to recover after launch
130
+ os.dup2(r1, _PYREMOTE_BOOTSTRAP_SRC_FD)
131
+
132
+ # Close remaining fd's
133
+ for f in [r0, w0, r1, w1]:
134
+ os.close(f)
135
+
136
+ # Save child pid to close after relaunch
137
+ os.environ[_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR] = str(cp)
138
+
139
+ # Save original argv0
140
+ os.environ[_PYREMOTE_BOOTSTRAP_ARGV0_VAR] = sys.executable
141
+
142
+ # Start repl reading stdin from r0
143
+ os.execl(sys.executable, sys.executable + (_PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT % (context_name,)))
144
+
145
+ else:
146
+ # Child process
147
+
148
+ # Write first ack
149
+ os.write(1, _PYREMOTE_BOOTSTRAP_ACK0)
150
+
151
+ # Write pid
152
+ os.write(1, struct.pack('<Q', pid))
153
+
154
+ # Read main src from stdin
155
+ main_z_len = struct.unpack('<I', os.read(0, 4))[0]
156
+ main_src = zlib.decompress(os.fdopen(0, 'rb').read(main_z_len))
157
+
158
+ # Write both copies of main src. Must write to w0 (parent stdin) before w1 (copy pipe) as pipe will likely fill
159
+ # and block and need to be drained by pyremote_bootstrap_finalize running in parent.
160
+ for w in [w0, w1]:
161
+ fp = os.fdopen(w, 'wb', 0)
162
+ fp.write(main_src)
163
+ fp.close()
164
+
165
+ # Write second ack
166
+ os.write(1, _PYREMOTE_BOOTSTRAP_ACK1)
167
+
168
+ # Exit child
169
+ sys.exit(0)
170
+
171
+
172
+ ##
173
+
174
+
175
+ def pyremote_build_bootstrap_cmd(context_name: str) -> str:
176
+ if any(c in context_name for c in '\'"'):
177
+ raise NameError(context_name)
178
+
179
+ bs_src = textwrap.dedent(inspect.getsource(_pyremote_bootstrap_main))
180
+
181
+ for gl in [
182
+ '_PYREMOTE_BOOTSTRAP_COMM_FD',
183
+ '_PYREMOTE_BOOTSTRAP_SRC_FD',
184
+
185
+ '_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR',
186
+ '_PYREMOTE_BOOTSTRAP_ARGV0_VAR',
187
+
188
+ '_PYREMOTE_BOOTSTRAP_ACK0',
189
+ '_PYREMOTE_BOOTSTRAP_ACK1',
190
+
191
+ '_PYREMOTE_BOOTSTRAP_PROC_TITLE_FMT',
192
+ ]:
193
+ bs_src = bs_src.replace(gl, repr(globals()[gl]))
194
+
195
+ bs_src = '\n'.join(
196
+ cl
197
+ for l in bs_src.splitlines()
198
+ if (cl := (l.split('#')[0]).rstrip())
199
+ if cl.strip()
200
+ )
201
+
202
+ bs_z = zlib.compress(bs_src.encode('utf-8'))
203
+ bs_z64 = base64.encodebytes(bs_z).replace(b'\n', b'')
204
+
205
+ def dq_repr(o: ta.Any) -> str:
206
+ return '"' + repr(o)[1:-1] + '"'
207
+
208
+ stmts = [
209
+ f'import {", ".join(_PYREMOTE_BOOTSTRAP_IMPORTS)}',
210
+ f'exec(zlib.decompress(base64.decodebytes({bs_z64!r})))',
211
+ f'_pyremote_bootstrap_main({context_name!r})',
212
+ ]
213
+
214
+ cmd = '; '.join(stmts)
215
+ return cmd
216
+
217
+
218
+ ##
219
+
220
+
221
+ @dc.dataclass(frozen=True)
222
+ class PyremoteEnvInfo:
223
+ sys_base_prefix: str
224
+ sys_byteorder: str
225
+ sys_defaultencoding: str
226
+ sys_exec_prefix: str
227
+ sys_executable: str
228
+ sys_implementation_name: str
229
+ sys_path: ta.List[str]
230
+ sys_platform: str
231
+ sys_prefix: str
232
+ sys_version: str
233
+ sys_version_info: ta.List[ta.Union[int, str]]
234
+
235
+ platform_architecture: ta.List[str]
236
+ platform_machine: str
237
+ platform_platform: str
238
+ platform_processor: str
239
+ platform_system: str
240
+ platform_release: str
241
+ platform_version: str
242
+
243
+ site_userbase: str
244
+
245
+ os_cwd: str
246
+ os_gid: int
247
+ os_loadavg: ta.List[float]
248
+ os_login: ta.Optional[str]
249
+ os_pgrp: int
250
+ os_pid: int
251
+ os_ppid: int
252
+ os_uid: int
253
+
254
+ pw_name: str
255
+ pw_uid: int
256
+ pw_gid: int
257
+ pw_gecos: str
258
+ pw_dir: str
259
+ pw_shell: str
260
+
261
+ env_path: ta.Optional[str]
262
+
263
+
264
+ def _get_pyremote_env_info() -> PyremoteEnvInfo:
265
+ os_uid = os.getuid()
266
+
267
+ pw = pwd.getpwuid(os_uid)
268
+
269
+ os_login: ta.Optional[str]
270
+ try:
271
+ os_login = os.getlogin()
272
+ except OSError:
273
+ os_login = None
274
+
275
+ return PyremoteEnvInfo(
276
+ sys_base_prefix=sys.base_prefix,
277
+ sys_byteorder=sys.byteorder,
278
+ sys_defaultencoding=sys.getdefaultencoding(),
279
+ sys_exec_prefix=sys.exec_prefix,
280
+ sys_executable=sys.executable,
281
+ sys_implementation_name=sys.implementation.name,
282
+ sys_path=sys.path,
283
+ sys_platform=sys.platform,
284
+ sys_prefix=sys.prefix,
285
+ sys_version=sys.version,
286
+ sys_version_info=list(sys.version_info),
287
+
288
+ platform_architecture=list(platform.architecture()),
289
+ platform_machine=platform.machine(),
290
+ platform_platform=platform.platform(),
291
+ platform_processor=platform.processor(),
292
+ platform_system=platform.system(),
293
+ platform_release=platform.release(),
294
+ platform_version=platform.version(),
295
+
296
+ site_userbase=site.getuserbase(),
297
+
298
+ os_cwd=os.getcwd(),
299
+ os_gid=os.getgid(),
300
+ os_loadavg=list(os.getloadavg()),
301
+ os_login=os_login,
302
+ os_pgrp=os.getpgrp(),
303
+ os_pid=os.getpid(),
304
+ os_ppid=os.getppid(),
305
+ os_uid=os_uid,
306
+
307
+ pw_name=pw.pw_name,
308
+ pw_uid=pw.pw_uid,
309
+ pw_gid=pw.pw_gid,
310
+ pw_gecos=pw.pw_gecos,
311
+ pw_dir=pw.pw_dir,
312
+ pw_shell=pw.pw_shell,
313
+
314
+ env_path=os.environ.get('PATH'),
315
+ )
316
+
317
+
318
+ ##
319
+
320
+
321
+ class PyremoteBootstrapDriver:
322
+ def __init__(self, main_src: str) -> None:
323
+ super().__init__()
324
+
325
+ self._main_src = main_src
326
+ self._main_z = zlib.compress(main_src.encode('utf-8'))
327
+
328
+ #
329
+
330
+ @dc.dataclass(frozen=True)
331
+ class Read:
332
+ sz: int
333
+
334
+ @dc.dataclass(frozen=True)
335
+ class Write:
336
+ d: bytes
337
+
338
+ class ProtocolError(Exception):
339
+ pass
340
+
341
+ @dc.dataclass(frozen=True)
342
+ class Result:
343
+ pid: int
344
+ env_info: PyremoteEnvInfo
345
+
346
+ def gen(self) -> ta.Generator[ta.Union[Read, Write], ta.Optional[bytes], Result]:
347
+ # Read first ack
348
+ yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK0)
349
+
350
+ # Read pid
351
+ d = yield from self._read(8)
352
+ pid = struct.unpack('<Q', d)[0]
353
+
354
+ # Write main src
355
+ yield from self._write(struct.pack('<I', len(self._main_z)))
356
+ yield from self._write(self._main_z)
357
+
358
+ # Read second and third ack
359
+ yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK1)
360
+ yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK2)
361
+
362
+ # Read env info
363
+ d = yield from self._read(4)
364
+ env_info_json_len = struct.unpack('<I', d)[0]
365
+ d = yield from self._read(env_info_json_len)
366
+ env_info_json = d.decode('utf-8')
367
+ env_info = PyremoteEnvInfo(**json.loads(env_info_json))
368
+
369
+ # Read fourth ack
370
+ yield from self._expect(_PYREMOTE_BOOTSTRAP_ACK3)
371
+
372
+ # Return
373
+ return self.Result(
374
+ pid=pid,
375
+ env_info=env_info,
376
+ )
377
+
378
+ def _read(self, sz: int) -> ta.Generator[Read, bytes, bytes]:
379
+ d = yield self.Read(sz)
380
+ if not isinstance(d, bytes):
381
+ raise self.ProtocolError(f'Expected bytes after read, got {d!r}')
382
+ if len(d) != sz:
383
+ raise self.ProtocolError(f'Read {len(d)} bytes, expected {sz}')
384
+ return d
385
+
386
+ def _expect(self, e: bytes) -> ta.Generator[Read, bytes, None]:
387
+ d = yield from self._read(len(e))
388
+ if d != e:
389
+ raise self.ProtocolError(f'Read {d!r}, expected {e!r}')
390
+
391
+ def _write(self, d: bytes) -> ta.Generator[Write, ta.Optional[bytes], None]:
392
+ i = yield self.Write(d)
393
+ if i is not None:
394
+ raise self.ProtocolError('Unexpected input after write')
395
+
396
+ #
397
+
398
+ def run(self, stdin: ta.IO, stdout: ta.IO) -> Result:
399
+ gen = self.gen()
400
+
401
+ gi: ta.Optional[bytes] = None
402
+ while True:
403
+ try:
404
+ if gi is not None:
405
+ go = gen.send(gi)
406
+ else:
407
+ go = next(gen)
408
+ except StopIteration as e:
409
+ return e.value
410
+
411
+ if isinstance(go, self.Read):
412
+ gi = stdout.read(go.sz)
413
+ elif isinstance(go, self.Write):
414
+ gi = None
415
+ stdin.write(go.d)
416
+ stdin.flush()
417
+ else:
418
+ raise TypeError(go)
419
+
420
+
421
+ ##
422
+
423
+
424
+ @dc.dataclass(frozen=True)
425
+ class PyremotePayloadRuntime:
426
+ input: ta.BinaryIO
427
+ main_src: str
428
+ env_info: PyremoteEnvInfo
429
+
430
+
431
+ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
432
+ # Restore original argv0
433
+ sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
434
+
435
+ # Read second copy of main src
436
+ r1 = os.fdopen(_PYREMOTE_BOOTSTRAP_SRC_FD, 'rb', 0)
437
+ main_src = r1.read().decode('utf-8')
438
+ r1.close()
439
+
440
+ # Reap boostrap child. Must be done after reading second copy of source because source may be too big to fit in a
441
+ # pipe at once.
442
+ os.waitpid(int(os.environ.pop(_PYREMOTE_BOOTSTRAP_CHILD_PID_VAR)), 0)
443
+
444
+ # Write third ack
445
+ os.write(1, _PYREMOTE_BOOTSTRAP_ACK2)
446
+
447
+ # Write env info
448
+ env_info = _get_pyremote_env_info()
449
+ env_info_json = json.dumps(dc.asdict(env_info), indent=None, separators=(',', ':')) # noqa
450
+ os.write(1, struct.pack('<I', len(env_info_json)))
451
+ os.write(1, env_info_json.encode('utf-8'))
452
+
453
+ # Write fourth ack
454
+ os.write(1, _PYREMOTE_BOOTSTRAP_ACK3)
455
+
456
+ # Return
457
+ return PyremotePayloadRuntime(
458
+ input=os.fdopen(_PYREMOTE_BOOTSTRAP_COMM_FD, 'rb', 0),
459
+ main_src=main_src,
460
+ env_info=env_info,
461
+ )
462
+
463
+
464
+ ########################################
465
+ # ../../../../omlish/lite/cached.py
466
+
467
+
468
+ class _cached_nullary: # noqa
469
+ def __init__(self, fn):
470
+ super().__init__()
471
+ self._fn = fn
472
+ self._value = self._missing = object()
473
+ functools.update_wrapper(self, fn)
474
+
475
+ def __call__(self, *args, **kwargs): # noqa
476
+ if self._value is self._missing:
477
+ self._value = self._fn()
478
+ return self._value
479
+
480
+ def __get__(self, instance, owner): # noqa
481
+ bound = instance.__dict__[self._fn.__name__] = self.__class__(self._fn.__get__(instance, owner))
482
+ return bound
483
+
484
+
485
+ def cached_nullary(fn): # ta.Callable[..., T]) -> ta.Callable[..., T]:
486
+ return _cached_nullary(fn)
487
+
488
+
489
+ ########################################
490
+ # ../../../../omlish/lite/check.py
491
+
492
+
493
+ def check_isinstance(v: ta.Any, spec: ta.Union[ta.Type[T], tuple]) -> T:
494
+ if not isinstance(v, spec):
495
+ raise TypeError(v)
496
+ return v
497
+
498
+
499
+ def check_not_isinstance(v: T, spec: ta.Union[type, tuple]) -> T:
500
+ if isinstance(v, spec):
501
+ raise TypeError(v)
502
+ return v
503
+
504
+
505
+ def check_none(v: T) -> None:
506
+ if v is not None:
507
+ raise ValueError(v)
508
+
509
+
510
+ def check_not_none(v: ta.Optional[T]) -> T:
511
+ if v is None:
512
+ raise ValueError
513
+ return v
514
+
515
+
516
+ def check_not(v: ta.Any) -> None:
517
+ if v:
518
+ raise ValueError(v)
519
+ return v
520
+
521
+
522
+ def check_non_empty_str(v: ta.Optional[str]) -> str:
523
+ if not v:
524
+ raise ValueError
525
+ return v
526
+
527
+
528
+ def check_state(v: bool, msg: str = 'Illegal state') -> None:
529
+ if not v:
530
+ raise ValueError(msg)
531
+
532
+
533
+ def check_equal(l: T, r: T) -> T:
534
+ if l != r:
535
+ raise ValueError(l, r)
536
+ return l
537
+
538
+
539
+ def check_not_equal(l: T, r: T) -> T:
540
+ if l == r:
541
+ raise ValueError(l, r)
542
+ return l
543
+
544
+
545
+ def check_is(l: T, r: T) -> T:
546
+ if l is not r:
547
+ raise ValueError(l, r)
548
+ return l
549
+
550
+
551
+ def check_is_not(l: T, r: ta.Any) -> T:
552
+ if l is r:
553
+ raise ValueError(l, r)
554
+ return l
555
+
556
+
557
+ def check_in(v: T, c: ta.Container[T]) -> T:
558
+ if v not in c:
559
+ raise ValueError(v, c)
560
+ return v
561
+
562
+
563
+ def check_not_in(v: T, c: ta.Container[T]) -> T:
564
+ if v in c:
565
+ raise ValueError(v, c)
566
+ return v
567
+
568
+
569
+ def check_single(vs: ta.Iterable[T]) -> T:
570
+ [v] = vs
571
+ return v
572
+
573
+
574
+ def check_empty(v: SizedT) -> SizedT:
575
+ if len(v):
576
+ raise ValueError(v)
577
+ return v
578
+
579
+
580
+ def check_non_empty(v: SizedT) -> SizedT:
581
+ if not len(v):
582
+ raise ValueError(v)
583
+ return v
584
+
585
+
586
+ ########################################
587
+ # ../../../../omlish/lite/json.py
588
+
589
+
590
+ ##
591
+
592
+
593
+ JSON_PRETTY_INDENT = 2
594
+
595
+ JSON_PRETTY_KWARGS: ta.Mapping[str, ta.Any] = dict(
596
+ indent=JSON_PRETTY_INDENT,
597
+ )
598
+
599
+ json_dump_pretty: ta.Callable[..., bytes] = functools.partial(json.dump, **JSON_PRETTY_KWARGS) # type: ignore
600
+ json_dumps_pretty: ta.Callable[..., str] = functools.partial(json.dumps, **JSON_PRETTY_KWARGS)
601
+
602
+
603
+ ##
604
+
605
+
606
+ JSON_COMPACT_SEPARATORS = (',', ':')
607
+
608
+ JSON_COMPACT_KWARGS: ta.Mapping[str, ta.Any] = dict(
609
+ indent=None,
610
+ separators=JSON_COMPACT_SEPARATORS,
611
+ )
612
+
613
+ json_dump_compact: ta.Callable[..., bytes] = functools.partial(json.dump, **JSON_COMPACT_KWARGS) # type: ignore
614
+ json_dumps_compact: ta.Callable[..., str] = functools.partial(json.dumps, **JSON_COMPACT_KWARGS)
615
+
616
+
617
+ ########################################
618
+ # ../../../../omlish/lite/reflect.py
619
+
620
+
621
+ _GENERIC_ALIAS_TYPES = (
622
+ ta._GenericAlias, # type: ignore # noqa
623
+ *([ta._SpecialGenericAlias] if hasattr(ta, '_SpecialGenericAlias') else []), # noqa
624
+ )
625
+
626
+
627
+ def is_generic_alias(obj, *, origin: ta.Any = None) -> bool:
628
+ return (
629
+ isinstance(obj, _GENERIC_ALIAS_TYPES) and
630
+ (origin is None or ta.get_origin(obj) is origin)
631
+ )
632
+
633
+
634
+ is_union_alias = functools.partial(is_generic_alias, origin=ta.Union)
635
+ is_callable_alias = functools.partial(is_generic_alias, origin=ta.Callable)
636
+
637
+
638
+ def is_optional_alias(spec: ta.Any) -> bool:
639
+ return (
640
+ isinstance(spec, _GENERIC_ALIAS_TYPES) and # noqa
641
+ ta.get_origin(spec) is ta.Union and
642
+ len(ta.get_args(spec)) == 2 and
643
+ any(a in (None, type(None)) for a in ta.get_args(spec))
644
+ )
645
+
646
+
647
+ def get_optional_alias_arg(spec: ta.Any) -> ta.Any:
648
+ [it] = [it for it in ta.get_args(spec) if it not in (None, type(None))]
649
+ return it
650
+
651
+
652
+ def is_new_type(spec: ta.Any) -> bool:
653
+ if isinstance(ta.NewType, type):
654
+ return isinstance(spec, ta.NewType)
655
+ else:
656
+ # Before https://github.com/python/cpython/commit/c2f33dfc83ab270412bf243fb21f724037effa1a
657
+ return isinstance(spec, types.FunctionType) and spec.__code__ is ta.NewType.__code__.co_consts[1] # type: ignore # noqa
658
+
659
+
660
+ def deep_subclasses(cls: ta.Type[T]) -> ta.Iterator[ta.Type[T]]:
661
+ seen = set()
662
+ todo = list(reversed(cls.__subclasses__()))
663
+ while todo:
664
+ cur = todo.pop()
665
+ if cur in seen:
666
+ continue
667
+ seen.add(cur)
668
+ yield cur
669
+ todo.extend(reversed(cur.__subclasses__()))
670
+
671
+
672
+ ########################################
673
+ # ../../../../omlish/lite/logs.py
674
+ """
675
+ TODO:
676
+ - translate json keys
677
+ - debug
678
+ """
679
+
680
+
681
+ log = logging.getLogger(__name__)
682
+
683
+
684
+ ##
685
+
686
+
687
+ class TidLogFilter(logging.Filter):
688
+
689
+ def filter(self, record):
690
+ record.tid = threading.get_native_id()
691
+ return True
692
+
693
+
694
+ ##
695
+
696
+
697
+ class JsonLogFormatter(logging.Formatter):
698
+
699
+ KEYS: ta.Mapping[str, bool] = {
700
+ 'name': False,
701
+ 'msg': False,
702
+ 'args': False,
703
+ 'levelname': False,
704
+ 'levelno': False,
705
+ 'pathname': False,
706
+ 'filename': False,
707
+ 'module': False,
708
+ 'exc_info': True,
709
+ 'exc_text': True,
710
+ 'stack_info': True,
711
+ 'lineno': False,
712
+ 'funcName': False,
713
+ 'created': False,
714
+ 'msecs': False,
715
+ 'relativeCreated': False,
716
+ 'thread': False,
717
+ 'threadName': False,
718
+ 'processName': False,
719
+ 'process': False,
720
+ }
721
+
722
+ def format(self, record: logging.LogRecord) -> str:
723
+ dct = {
724
+ k: v
725
+ for k, o in self.KEYS.items()
726
+ for v in [getattr(record, k)]
727
+ if not (o and v is None)
728
+ }
729
+ return json_dumps_compact(dct)
730
+
731
+
732
+ ##
733
+
734
+
735
+ STANDARD_LOG_FORMAT_PARTS = [
736
+ ('asctime', '%(asctime)-15s'),
737
+ ('process', 'pid=%(process)-6s'),
738
+ ('thread', 'tid=%(thread)x'),
739
+ ('levelname', '%(levelname)s'),
740
+ ('name', '%(name)s'),
741
+ ('separator', '::'),
742
+ ('message', '%(message)s'),
743
+ ]
744
+
745
+
746
+ class StandardLogFormatter(logging.Formatter):
747
+
748
+ @staticmethod
749
+ def build_log_format(parts: ta.Iterable[ta.Tuple[str, str]]) -> str:
750
+ return ' '.join(v for k, v in parts)
751
+
752
+ converter = datetime.datetime.fromtimestamp # type: ignore
753
+
754
+ def formatTime(self, record, datefmt=None):
755
+ ct = self.converter(record.created) # type: ignore
756
+ if datefmt:
757
+ return ct.strftime(datefmt) # noqa
758
+ else:
759
+ t = ct.strftime('%Y-%m-%d %H:%M:%S')
760
+ return '%s.%03d' % (t, record.msecs) # noqa
761
+
762
+
763
+ ##
764
+
765
+
766
+ class ProxyLogFilterer(logging.Filterer):
767
+ def __init__(self, underlying: logging.Filterer) -> None: # noqa
768
+ self._underlying = underlying
769
+
770
+ @property
771
+ def underlying(self) -> logging.Filterer:
772
+ return self._underlying
773
+
774
+ @property
775
+ def filters(self):
776
+ return self._underlying.filters
777
+
778
+ @filters.setter
779
+ def filters(self, filters):
780
+ self._underlying.filters = filters
781
+
782
+ def addFilter(self, filter): # noqa
783
+ self._underlying.addFilter(filter)
784
+
785
+ def removeFilter(self, filter): # noqa
786
+ self._underlying.removeFilter(filter)
787
+
788
+ def filter(self, record):
789
+ return self._underlying.filter(record)
790
+
791
+
792
+ class ProxyLogHandler(ProxyLogFilterer, logging.Handler):
793
+ def __init__(self, underlying: logging.Handler) -> None: # noqa
794
+ ProxyLogFilterer.__init__(self, underlying)
795
+
796
+ _underlying: logging.Handler
797
+
798
+ @property
799
+ def underlying(self) -> logging.Handler:
800
+ return self._underlying
801
+
802
+ def get_name(self):
803
+ return self._underlying.get_name()
804
+
805
+ def set_name(self, name):
806
+ self._underlying.set_name(name)
807
+
808
+ @property
809
+ def name(self):
810
+ return self._underlying.name
811
+
812
+ @property
813
+ def level(self):
814
+ return self._underlying.level
815
+
816
+ @level.setter
817
+ def level(self, level):
818
+ self._underlying.level = level
819
+
820
+ @property
821
+ def formatter(self):
822
+ return self._underlying.formatter
823
+
824
+ @formatter.setter
825
+ def formatter(self, formatter):
826
+ self._underlying.formatter = formatter
827
+
828
+ def createLock(self):
829
+ self._underlying.createLock()
830
+
831
+ def acquire(self):
832
+ self._underlying.acquire()
833
+
834
+ def release(self):
835
+ self._underlying.release()
836
+
837
+ def setLevel(self, level):
838
+ self._underlying.setLevel(level)
839
+
840
+ def format(self, record):
841
+ return self._underlying.format(record)
842
+
843
+ def emit(self, record):
844
+ self._underlying.emit(record)
845
+
846
+ def handle(self, record):
847
+ return self._underlying.handle(record)
848
+
849
+ def setFormatter(self, fmt):
850
+ self._underlying.setFormatter(fmt)
851
+
852
+ def flush(self):
853
+ self._underlying.flush()
854
+
855
+ def close(self):
856
+ self._underlying.close()
857
+
858
+ def handleError(self, record):
859
+ self._underlying.handleError(record)
860
+
861
+
862
+ ##
863
+
864
+
865
+ class StandardLogHandler(ProxyLogHandler):
866
+ pass
867
+
868
+
869
+ ##
870
+
871
+
872
+ @contextlib.contextmanager
873
+ def _locking_logging_module_lock() -> ta.Iterator[None]:
874
+ if hasattr(logging, '_acquireLock'):
875
+ logging._acquireLock() # noqa
876
+ try:
877
+ yield
878
+ finally:
879
+ logging._releaseLock() # type: ignore # noqa
880
+
881
+ elif hasattr(logging, '_lock'):
882
+ # https://github.com/python/cpython/commit/74723e11109a320e628898817ab449b3dad9ee96
883
+ with logging._lock: # noqa
884
+ yield
885
+
886
+ else:
887
+ raise Exception("Can't find lock in logging module")
888
+
889
+
890
+ def configure_standard_logging(
891
+ level: ta.Union[int, str] = logging.INFO,
892
+ *,
893
+ json: bool = False,
894
+ target: ta.Optional[logging.Logger] = None,
895
+ force: bool = False,
896
+ handler_factory: ta.Optional[ta.Callable[[], logging.Handler]] = None,
897
+ ) -> ta.Optional[StandardLogHandler]:
898
+ with _locking_logging_module_lock():
899
+ if target is None:
900
+ target = logging.root
901
+
902
+ #
903
+
904
+ if not force:
905
+ if any(isinstance(h, StandardLogHandler) for h in list(target.handlers)):
906
+ return None
907
+
908
+ #
909
+
910
+ if handler_factory is not None:
911
+ handler = handler_factory()
912
+ else:
913
+ handler = logging.StreamHandler()
914
+
915
+ #
916
+
917
+ formatter: logging.Formatter
918
+ if json:
919
+ formatter = JsonLogFormatter()
920
+ else:
921
+ formatter = StandardLogFormatter(StandardLogFormatter.build_log_format(STANDARD_LOG_FORMAT_PARTS))
922
+ handler.setFormatter(formatter)
923
+
924
+ #
925
+
926
+ handler.addFilter(TidLogFilter())
927
+
928
+ #
929
+
930
+ target.addHandler(handler)
931
+
932
+ #
933
+
934
+ if level is not None:
935
+ target.setLevel(level)
936
+
937
+ #
938
+
939
+ return StandardLogHandler(handler)
940
+
941
+
942
+ ########################################
943
+ # ../../../../omlish/lite/marshal.py
944
+ """
945
+ TODO:
946
+ - pickle stdlib objs? have to pin to 3.8 pickle protocol, will be cross-version
947
+ - nonstrict toggle
948
+ """
949
+
950
+
951
+ ##
952
+
953
+
954
+ class ObjMarshaler(abc.ABC):
955
+ @abc.abstractmethod
956
+ def marshal(self, o: ta.Any) -> ta.Any:
957
+ raise NotImplementedError
958
+
959
+ @abc.abstractmethod
960
+ def unmarshal(self, o: ta.Any) -> ta.Any:
961
+ raise NotImplementedError
962
+
963
+
964
+ class NopObjMarshaler(ObjMarshaler):
965
+ def marshal(self, o: ta.Any) -> ta.Any:
966
+ return o
967
+
968
+ def unmarshal(self, o: ta.Any) -> ta.Any:
969
+ return o
970
+
971
+
972
+ @dc.dataclass()
973
+ class ProxyObjMarshaler(ObjMarshaler):
974
+ m: ta.Optional[ObjMarshaler] = None
975
+
976
+ def marshal(self, o: ta.Any) -> ta.Any:
977
+ return check_not_none(self.m).marshal(o)
978
+
979
+ def unmarshal(self, o: ta.Any) -> ta.Any:
980
+ return check_not_none(self.m).unmarshal(o)
981
+
982
+
983
+ @dc.dataclass(frozen=True)
984
+ class CastObjMarshaler(ObjMarshaler):
985
+ ty: type
986
+
987
+ def marshal(self, o: ta.Any) -> ta.Any:
988
+ return o
989
+
990
+ def unmarshal(self, o: ta.Any) -> ta.Any:
991
+ return self.ty(o)
992
+
993
+
994
+ class DynamicObjMarshaler(ObjMarshaler):
995
+ def marshal(self, o: ta.Any) -> ta.Any:
996
+ return marshal_obj(o)
997
+
998
+ def unmarshal(self, o: ta.Any) -> ta.Any:
999
+ return o
1000
+
1001
+
1002
+ @dc.dataclass(frozen=True)
1003
+ class Base64ObjMarshaler(ObjMarshaler):
1004
+ ty: type
1005
+
1006
+ def marshal(self, o: ta.Any) -> ta.Any:
1007
+ return base64.b64encode(o).decode('ascii')
1008
+
1009
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1010
+ return self.ty(base64.b64decode(o))
1011
+
1012
+
1013
+ @dc.dataclass(frozen=True)
1014
+ class EnumObjMarshaler(ObjMarshaler):
1015
+ ty: type
1016
+
1017
+ def marshal(self, o: ta.Any) -> ta.Any:
1018
+ return o.name
1019
+
1020
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1021
+ return self.ty.__members__[o] # type: ignore
1022
+
1023
+
1024
+ @dc.dataclass(frozen=True)
1025
+ class OptionalObjMarshaler(ObjMarshaler):
1026
+ item: ObjMarshaler
1027
+
1028
+ def marshal(self, o: ta.Any) -> ta.Any:
1029
+ if o is None:
1030
+ return None
1031
+ return self.item.marshal(o)
1032
+
1033
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1034
+ if o is None:
1035
+ return None
1036
+ return self.item.unmarshal(o)
1037
+
1038
+
1039
+ @dc.dataclass(frozen=True)
1040
+ class MappingObjMarshaler(ObjMarshaler):
1041
+ ty: type
1042
+ km: ObjMarshaler
1043
+ vm: ObjMarshaler
1044
+
1045
+ def marshal(self, o: ta.Any) -> ta.Any:
1046
+ return {self.km.marshal(k): self.vm.marshal(v) for k, v in o.items()}
1047
+
1048
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1049
+ return self.ty((self.km.unmarshal(k), self.vm.unmarshal(v)) for k, v in o.items())
1050
+
1051
+
1052
+ @dc.dataclass(frozen=True)
1053
+ class IterableObjMarshaler(ObjMarshaler):
1054
+ ty: type
1055
+ item: ObjMarshaler
1056
+
1057
+ def marshal(self, o: ta.Any) -> ta.Any:
1058
+ return [self.item.marshal(e) for e in o]
1059
+
1060
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1061
+ return self.ty(self.item.unmarshal(e) for e in o)
1062
+
1063
+
1064
+ @dc.dataclass(frozen=True)
1065
+ class DataclassObjMarshaler(ObjMarshaler):
1066
+ ty: type
1067
+ fs: ta.Mapping[str, ObjMarshaler]
1068
+ nonstrict: bool = False
1069
+
1070
+ def marshal(self, o: ta.Any) -> ta.Any:
1071
+ return {k: m.marshal(getattr(o, k)) for k, m in self.fs.items()}
1072
+
1073
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1074
+ return self.ty(**{k: self.fs[k].unmarshal(v) for k, v in o.items() if not self.nonstrict or k in self.fs})
1075
+
1076
+
1077
+ @dc.dataclass(frozen=True)
1078
+ class PolymorphicObjMarshaler(ObjMarshaler):
1079
+ class Impl(ta.NamedTuple):
1080
+ ty: type
1081
+ tag: str
1082
+ m: ObjMarshaler
1083
+
1084
+ impls_by_ty: ta.Mapping[type, Impl]
1085
+ impls_by_tag: ta.Mapping[str, Impl]
1086
+
1087
+ @classmethod
1088
+ def of(cls, impls: ta.Iterable[Impl]) -> 'PolymorphicObjMarshaler':
1089
+ return cls(
1090
+ {i.ty: i for i in impls},
1091
+ {i.tag: i for i in impls},
1092
+ )
1093
+
1094
+ def marshal(self, o: ta.Any) -> ta.Any:
1095
+ impl = self.impls_by_ty[type(o)]
1096
+ return {impl.tag: impl.m.marshal(o)}
1097
+
1098
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1099
+ [(t, v)] = o.items()
1100
+ impl = self.impls_by_tag[t]
1101
+ return impl.m.unmarshal(v)
1102
+
1103
+
1104
+ @dc.dataclass(frozen=True)
1105
+ class DatetimeObjMarshaler(ObjMarshaler):
1106
+ ty: type
1107
+
1108
+ def marshal(self, o: ta.Any) -> ta.Any:
1109
+ return o.isoformat()
1110
+
1111
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1112
+ return self.ty.fromisoformat(o) # type: ignore
1113
+
1114
+
1115
+ class DecimalObjMarshaler(ObjMarshaler):
1116
+ def marshal(self, o: ta.Any) -> ta.Any:
1117
+ return str(check_isinstance(o, decimal.Decimal))
1118
+
1119
+ def unmarshal(self, v: ta.Any) -> ta.Any:
1120
+ return decimal.Decimal(check_isinstance(v, str))
1121
+
1122
+
1123
+ class FractionObjMarshaler(ObjMarshaler):
1124
+ def marshal(self, o: ta.Any) -> ta.Any:
1125
+ fr = check_isinstance(o, fractions.Fraction)
1126
+ return [fr.numerator, fr.denominator]
1127
+
1128
+ def unmarshal(self, v: ta.Any) -> ta.Any:
1129
+ num, denom = check_isinstance(v, list)
1130
+ return fractions.Fraction(num, denom)
1131
+
1132
+
1133
+ class UuidObjMarshaler(ObjMarshaler):
1134
+ def marshal(self, o: ta.Any) -> ta.Any:
1135
+ return str(o)
1136
+
1137
+ def unmarshal(self, o: ta.Any) -> ta.Any:
1138
+ return uuid.UUID(o)
1139
+
1140
+
1141
+ ##
1142
+
1143
+
1144
+ _DEFAULT_OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = {
1145
+ **{t: NopObjMarshaler() for t in (type(None),)},
1146
+ **{t: CastObjMarshaler(t) for t in (int, float, str, bool)},
1147
+ **{t: Base64ObjMarshaler(t) for t in (bytes, bytearray)},
1148
+ **{t: IterableObjMarshaler(t, DynamicObjMarshaler()) for t in (list, tuple, set, frozenset)},
1149
+ **{t: MappingObjMarshaler(t, DynamicObjMarshaler(), DynamicObjMarshaler()) for t in (dict,)},
1150
+
1151
+ ta.Any: DynamicObjMarshaler(),
1152
+
1153
+ **{t: DatetimeObjMarshaler(t) for t in (datetime.date, datetime.time, datetime.datetime)},
1154
+ decimal.Decimal: DecimalObjMarshaler(),
1155
+ fractions.Fraction: FractionObjMarshaler(),
1156
+ uuid.UUID: UuidObjMarshaler(),
1157
+ }
1158
+
1159
+ _OBJ_MARSHALER_GENERIC_MAPPING_TYPES: ta.Dict[ta.Any, type] = {
1160
+ **{t: t for t in (dict,)},
1161
+ **{t: dict for t in (collections.abc.Mapping, collections.abc.MutableMapping)},
1162
+ }
1163
+
1164
+ _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES: ta.Dict[ta.Any, type] = {
1165
+ **{t: t for t in (list, tuple, set, frozenset)},
1166
+ collections.abc.Set: frozenset,
1167
+ collections.abc.MutableSet: set,
1168
+ collections.abc.Sequence: tuple,
1169
+ collections.abc.MutableSequence: list,
1170
+ }
1171
+
1172
+
1173
+ def _make_obj_marshaler(
1174
+ ty: ta.Any,
1175
+ rec: ta.Callable[[ta.Any], ObjMarshaler],
1176
+ *,
1177
+ nonstrict_dataclasses: bool = False,
1178
+ ) -> ObjMarshaler:
1179
+ if isinstance(ty, type):
1180
+ if abc.ABC in ty.__bases__:
1181
+ return PolymorphicObjMarshaler.of([ # type: ignore
1182
+ PolymorphicObjMarshaler.Impl(
1183
+ ity,
1184
+ ity.__qualname__,
1185
+ rec(ity),
1186
+ )
1187
+ for ity in deep_subclasses(ty)
1188
+ if abc.ABC not in ity.__bases__
1189
+ ])
1190
+
1191
+ if issubclass(ty, enum.Enum):
1192
+ return EnumObjMarshaler(ty)
1193
+
1194
+ if dc.is_dataclass(ty):
1195
+ return DataclassObjMarshaler(
1196
+ ty,
1197
+ {f.name: rec(f.type) for f in dc.fields(ty)},
1198
+ nonstrict=nonstrict_dataclasses,
1199
+ )
1200
+
1201
+ if is_generic_alias(ty):
1202
+ try:
1203
+ mt = _OBJ_MARSHALER_GENERIC_MAPPING_TYPES[ta.get_origin(ty)]
1204
+ except KeyError:
1205
+ pass
1206
+ else:
1207
+ k, v = ta.get_args(ty)
1208
+ return MappingObjMarshaler(mt, rec(k), rec(v))
1209
+
1210
+ try:
1211
+ st = _OBJ_MARSHALER_GENERIC_ITERABLE_TYPES[ta.get_origin(ty)]
1212
+ except KeyError:
1213
+ pass
1214
+ else:
1215
+ [e] = ta.get_args(ty)
1216
+ return IterableObjMarshaler(st, rec(e))
1217
+
1218
+ if is_union_alias(ty):
1219
+ return OptionalObjMarshaler(rec(get_optional_alias_arg(ty)))
1220
+
1221
+ raise TypeError(ty)
1222
+
1223
+
1224
+ ##
1225
+
1226
+
1227
+ _OBJ_MARSHALERS_LOCK = threading.RLock()
1228
+
1229
+ _OBJ_MARSHALERS: ta.Dict[ta.Any, ObjMarshaler] = dict(_DEFAULT_OBJ_MARSHALERS)
1230
+
1231
+ _OBJ_MARSHALER_PROXIES: ta.Dict[ta.Any, ProxyObjMarshaler] = {}
1232
+
1233
+
1234
+ def register_opj_marshaler(ty: ta.Any, m: ObjMarshaler) -> None:
1235
+ with _OBJ_MARSHALERS_LOCK:
1236
+ if ty in _OBJ_MARSHALERS:
1237
+ raise KeyError(ty)
1238
+ _OBJ_MARSHALERS[ty] = m
1239
+
1240
+
1241
+ def get_obj_marshaler(
1242
+ ty: ta.Any,
1243
+ *,
1244
+ no_cache: bool = False,
1245
+ **kwargs: ta.Any,
1246
+ ) -> ObjMarshaler:
1247
+ with _OBJ_MARSHALERS_LOCK:
1248
+ if not no_cache:
1249
+ try:
1250
+ return _OBJ_MARSHALERS[ty]
1251
+ except KeyError:
1252
+ pass
1253
+
1254
+ try:
1255
+ return _OBJ_MARSHALER_PROXIES[ty]
1256
+ except KeyError:
1257
+ pass
1258
+
1259
+ rec = functools.partial(
1260
+ get_obj_marshaler,
1261
+ no_cache=no_cache,
1262
+ **kwargs,
1263
+ )
1264
+
1265
+ p = ProxyObjMarshaler()
1266
+ _OBJ_MARSHALER_PROXIES[ty] = p
1267
+ try:
1268
+ m = _make_obj_marshaler(ty, rec, **kwargs)
1269
+ finally:
1270
+ del _OBJ_MARSHALER_PROXIES[ty]
1271
+ p.m = m
1272
+
1273
+ if not no_cache:
1274
+ _OBJ_MARSHALERS[ty] = m
1275
+ return m
1276
+
1277
+
1278
+ ##
1279
+
1280
+
1281
+ def marshal_obj(o: ta.Any, ty: ta.Any = None) -> ta.Any:
1282
+ return get_obj_marshaler(ty if ty is not None else type(o)).marshal(o)
1283
+
1284
+
1285
+ def unmarshal_obj(o: ta.Any, ty: ta.Union[ta.Type[T], ta.Any]) -> T:
1286
+ return get_obj_marshaler(ty).unmarshal(o)
1287
+
1288
+
1289
+ ########################################
1290
+ # ../../../../omlish/lite/runtime.py
1291
+
1292
+
1293
+ @cached_nullary
1294
+ def is_debugger_attached() -> bool:
1295
+ return any(frame[1].endswith('pydevd.py') for frame in inspect.stack())
1296
+
1297
+
1298
+ REQUIRED_PYTHON_VERSION = (3, 8)
1299
+
1300
+
1301
+ def check_runtime_version() -> None:
1302
+ if sys.version_info < REQUIRED_PYTHON_VERSION:
1303
+ raise OSError(f'Requires python {REQUIRED_PYTHON_VERSION}, got {sys.version_info} from {sys.executable}') # noqa
1304
+
1305
+
1306
+ ########################################
1307
+ # ../../../../omlish/lite/subprocesses.py
1308
+
1309
+
1310
+ ##
1311
+
1312
+
1313
+ _SUBPROCESS_SHELL_WRAP_EXECS = False
1314
+
1315
+
1316
+ def subprocess_shell_wrap_exec(*args: str) -> ta.Tuple[str, ...]:
1317
+ return ('sh', '-c', ' '.join(map(shlex.quote, args)))
1318
+
1319
+
1320
+ def subprocess_maybe_shell_wrap_exec(*args: str) -> ta.Tuple[str, ...]:
1321
+ if _SUBPROCESS_SHELL_WRAP_EXECS or is_debugger_attached():
1322
+ return subprocess_shell_wrap_exec(*args)
1323
+ else:
1324
+ return args
1325
+
1326
+
1327
+ def _prepare_subprocess_invocation(
1328
+ *args: str,
1329
+ env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
1330
+ extra_env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
1331
+ quiet: bool = False,
1332
+ shell: bool = False,
1333
+ **kwargs: ta.Any,
1334
+ ) -> ta.Tuple[ta.Tuple[ta.Any, ...], ta.Dict[str, ta.Any]]:
1335
+ log.debug(args)
1336
+ if extra_env:
1337
+ log.debug(extra_env)
1338
+
1339
+ if extra_env:
1340
+ env = {**(env if env is not None else os.environ), **extra_env}
1341
+
1342
+ if quiet and 'stderr' not in kwargs:
1343
+ if not log.isEnabledFor(logging.DEBUG):
1344
+ kwargs['stderr'] = subprocess.DEVNULL
1345
+
1346
+ if not shell:
1347
+ args = subprocess_maybe_shell_wrap_exec(*args)
1348
+
1349
+ return args, dict(
1350
+ env=env,
1351
+ shell=shell,
1352
+ **kwargs,
1353
+ )
1354
+
1355
+
1356
+ def subprocess_check_call(*args: str, stdout=sys.stderr, **kwargs: ta.Any) -> None:
1357
+ args, kwargs = _prepare_subprocess_invocation(*args, stdout=stdout, **kwargs)
1358
+ return subprocess.check_call(args, **kwargs) # type: ignore
1359
+
1360
+
1361
+ def subprocess_check_output(*args: str, **kwargs: ta.Any) -> bytes:
1362
+ args, kwargs = _prepare_subprocess_invocation(*args, **kwargs)
1363
+ return subprocess.check_output(args, **kwargs)
1364
+
1365
+
1366
+ def subprocess_check_output_str(*args: str, **kwargs: ta.Any) -> str:
1367
+ return subprocess_check_output(*args, **kwargs).decode().strip()
1368
+
1369
+
1370
+ ##
1371
+
1372
+
1373
+ DEFAULT_SUBPROCESS_TRY_EXCEPTIONS: ta.Tuple[ta.Type[Exception], ...] = (
1374
+ FileNotFoundError,
1375
+ subprocess.CalledProcessError,
1376
+ )
1377
+
1378
+
1379
+ def subprocess_try_call(
1380
+ *args: str,
1381
+ try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
1382
+ **kwargs: ta.Any,
1383
+ ) -> bool:
1384
+ try:
1385
+ subprocess_check_call(*args, **kwargs)
1386
+ except try_exceptions as e: # noqa
1387
+ if log.isEnabledFor(logging.DEBUG):
1388
+ log.exception('command failed')
1389
+ return False
1390
+ else:
1391
+ return True
1392
+
1393
+
1394
+ def subprocess_try_output(
1395
+ *args: str,
1396
+ try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
1397
+ **kwargs: ta.Any,
1398
+ ) -> ta.Optional[bytes]:
1399
+ try:
1400
+ return subprocess_check_output(*args, **kwargs)
1401
+ except try_exceptions as e: # noqa
1402
+ if log.isEnabledFor(logging.DEBUG):
1403
+ log.exception('command failed')
1404
+ return None
1405
+
1406
+
1407
+ def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
1408
+ out = subprocess_try_output(*args, **kwargs)
1409
+ return out.decode().strip() if out is not None else None
1410
+
1411
+
1412
+ ##
1413
+
1414
+
1415
+ def subprocess_close(
1416
+ proc: subprocess.Popen,
1417
+ timeout: ta.Optional[float] = None,
1418
+ ) -> None:
1419
+ # TODO: terminate, sleep, kill
1420
+ if proc.stdout:
1421
+ proc.stdout.close()
1422
+ if proc.stderr:
1423
+ proc.stderr.close()
1424
+ if proc.stdin:
1425
+ proc.stdin.close()
1426
+
1427
+ proc.wait(timeout)
1428
+
1429
+
1430
+ ########################################
1431
+ # ../commands/subprocess.py
1432
+
1433
+
1434
+ ##
1435
+
1436
+
1437
+ class SubprocessCommand(Command['SubprocessCommand.Input', 'SubprocessCommand.Output']):
1438
+ @dc.dataclass(frozen=True)
1439
+ class Input(Command.Input):
1440
+ args: ta.Sequence[str]
1441
+
1442
+ shell: bool = False
1443
+ cwd: ta.Optional[str] = None
1444
+ env: ta.Optional[ta.Mapping[str, str]] = None
1445
+
1446
+ capture_stdout: bool = False
1447
+ capture_stderr: bool = False
1448
+
1449
+ input: ta.Optional[bytes] = None
1450
+ timeout: ta.Optional[float] = None
1451
+
1452
+ def __post_init__(self) -> None:
1453
+ if isinstance(self.args, str):
1454
+ raise TypeError(self.args)
1455
+
1456
+ @dc.dataclass(frozen=True)
1457
+ class Output(Command.Output):
1458
+ rc: int
1459
+ pid: int
1460
+
1461
+ elapsed_s: float
1462
+
1463
+ stdout: ta.Optional[bytes] = None
1464
+ stderr: ta.Optional[bytes] = None
1465
+
1466
+ def _execute(self, inp: Input) -> Output:
1467
+ proc = subprocess.Popen(
1468
+ subprocess_maybe_shell_wrap_exec(*inp.args),
1469
+
1470
+ shell=inp.shell,
1471
+ cwd=inp.cwd,
1472
+ env={**os.environ, **(inp.env or {})},
1473
+
1474
+ stdin=subprocess.PIPE if inp.input is not None else None,
1475
+ stdout=subprocess.PIPE if inp.capture_stdout else None,
1476
+ stderr=subprocess.PIPE if inp.capture_stderr else None,
1477
+ )
1478
+
1479
+ start_time = time.time()
1480
+ stdout, stderr = proc.communicate(
1481
+ input=inp.input,
1482
+ timeout=inp.timeout,
1483
+ )
1484
+ end_time = time.time()
1485
+
1486
+ return SubprocessCommand.Output(
1487
+ rc=proc.returncode,
1488
+ pid=proc.pid,
1489
+
1490
+ elapsed_s=end_time - start_time,
1491
+
1492
+ stdout=stdout, # noqa
1493
+ stderr=stderr, # noqa
1494
+ )
1495
+
1496
+
1497
+ ########################################
1498
+ # main.py
1499
+
1500
+
1501
+ ##
1502
+
1503
+
1504
+ def _send_obj(f: ta.IO, o: ta.Any) -> None:
1505
+ j = json_dumps_compact(marshal_obj(o))
1506
+ d = j.encode('utf-8')
1507
+
1508
+ f.write(struct.pack('<I', len(d)))
1509
+ f.write(d)
1510
+ f.flush()
1511
+
1512
+
1513
+ def _recv_obj(f: ta.IO, ty: type) -> ta.Any:
1514
+ d = f.read(4)
1515
+ if not d:
1516
+ return None
1517
+ if len(d) != 4:
1518
+ raise Exception
1519
+
1520
+ sz = struct.unpack('<I', d)[0]
1521
+ d = f.read(sz)
1522
+ if not d:
1523
+ raise Exception
1524
+
1525
+ j = json.loads(d.decode('utf-8'))
1526
+ return unmarshal_obj(j, ty)
1527
+
1528
+
1529
+ def _remote_main() -> None:
1530
+ rt = pyremote_bootstrap_finalize() # noqa
1531
+
1532
+ while True:
1533
+ i = _recv_obj(rt.input, SubprocessCommand.Input)
1534
+ if i is None:
1535
+ break
1536
+
1537
+ o = SubprocessCommand()._execute(i) # noqa
1538
+
1539
+ _send_obj(sys.stdout.buffer, o)
1540
+
1541
+
1542
+ def _main() -> None:
1543
+ import argparse
1544
+
1545
+ parser = argparse.ArgumentParser()
1546
+
1547
+ parser.add_argument('--ssh')
1548
+ parser.add_argument('--python', default='python3')
1549
+ parser.add_argument('--_amalg-file')
1550
+
1551
+ args = parser.parse_args()
1552
+
1553
+ #
1554
+
1555
+ self_src = inspect.getsource(sys.modules[__name__])
1556
+ self_src_lines = self_src.splitlines()
1557
+ for l in self_src_lines:
1558
+ if l.startswith('# @omlish-amalg-output '):
1559
+ is_self_amalg = True
1560
+ break
1561
+ else:
1562
+ is_self_amalg = False
1563
+
1564
+ if is_self_amalg:
1565
+ amalg_src = self_src
1566
+ else:
1567
+ amalg_file = args._amalg_file # noqa
1568
+ if amalg_file is None:
1569
+ amalg_file = os.path.join(os.path.dirname(__file__), '_manage.py')
1570
+ with open(amalg_file) as f:
1571
+ amalg_src = f.read()
1572
+
1573
+ #
1574
+
1575
+ remote_src = '\n\n'.join([
1576
+ '__name__ = "__remote__"',
1577
+ amalg_src,
1578
+ '_remote_main()',
1579
+ ])
1580
+
1581
+ #
1582
+
1583
+ bs_src = pyremote_build_bootstrap_cmd(__package__ or 'manage')
1584
+
1585
+ if args.ssh is not None:
1586
+ sh_src = ' '.join([args.python, '-c', shlex.quote(bs_src)])
1587
+ sh_cmd = f'{args.ssh} {shlex.quote(sh_src)}'
1588
+ cmd = [sh_cmd]
1589
+ shell = True
1590
+ else:
1591
+ cmd = [args.python, '-c', bs_src]
1592
+ shell = False
1593
+
1594
+ proc = subprocess.Popen(
1595
+ subprocess_maybe_shell_wrap_exec(*cmd),
1596
+ shell=shell,
1597
+ stdin=subprocess.PIPE,
1598
+ stdout=subprocess.PIPE,
1599
+ )
1600
+
1601
+ stdin = check_not_none(proc.stdin)
1602
+ stdout = check_not_none(proc.stdout)
1603
+
1604
+ res = PyremoteBootstrapDriver(remote_src).run(stdin, stdout)
1605
+ print(res)
1606
+
1607
+ #
1608
+
1609
+ for ci in [
1610
+ SubprocessCommand.Input(
1611
+ args=['python3', '-'],
1612
+ input=b'print(1)\n',
1613
+ capture_stdout=True,
1614
+ ),
1615
+ SubprocessCommand.Input(
1616
+ args=['uname'],
1617
+ capture_stdout=True,
1618
+ ),
1619
+ ]:
1620
+ _send_obj(stdin, ci)
1621
+
1622
+ o = _recv_obj(stdout, SubprocessCommand.Output)
1623
+
1624
+ print(o)
1625
+
1626
+ try:
1627
+ stdin.close()
1628
+ except BrokenPipeError:
1629
+ pass
1630
+
1631
+ proc.wait()
1632
+
1633
+
1634
+ if __name__ == '__main__':
1635
+ _main()