ominfra 0.0.0.dev147__py3-none-any.whl → 0.0.0.dev148__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
ominfra/scripts/manage.py CHANGED
@@ -3,16 +3,20 @@
3
3
  # @omlish-lite
4
4
  # @omlish-script
5
5
  # @omlish-amalg-output ../manage/main.py
6
- # ruff: noqa: N802 UP006 UP007 UP036
6
+ # ruff: noqa: N802 TC003 UP006 UP007 UP036
7
7
  """
8
8
  manage.py -s 'docker run -i python:3.12'
9
9
  manage.py -s 'ssh -i /foo/bar.pem foo@bar.baz' -q --python=python3.8
10
10
  """
11
11
  import abc
12
+ import asyncio
13
+ import asyncio.base_subprocess
14
+ import asyncio.subprocess
12
15
  import base64
13
16
  import collections
14
17
  import collections.abc
15
18
  import contextlib
19
+ import ctypes as ct
16
20
  import dataclasses as dc
17
21
  import datetime
18
22
  import decimal
@@ -30,6 +34,7 @@ import pwd
30
34
  import re
31
35
  import shlex
32
36
  import shutil
37
+ import signal
33
38
  import site
34
39
  import struct
35
40
  import subprocess
@@ -62,6 +67,9 @@ VersionCmpLocalType = ta.Union['NegativeInfinityVersionType', _VersionCmpLocalTy
62
67
  VersionCmpKey = ta.Tuple[int, ta.Tuple[int, ...], VersionCmpPrePostDevType, VersionCmpPrePostDevType, VersionCmpPrePostDevType, VersionCmpLocalType] # noqa
63
68
  VersionComparisonMethod = ta.Callable[[VersionCmpKey, VersionCmpKey], bool]
64
69
 
70
+ # ../../omlish/lite/asyncio/asyncio.py
71
+ AwaitableT = ta.TypeVar('AwaitableT', bound=ta.Awaitable)
72
+
65
73
  # ../../omlish/lite/cached.py
66
74
  T = ta.TypeVar('T')
67
75
  CallableT = ta.TypeVar('CallableT', bound=ta.Callable)
@@ -822,10 +830,8 @@ def pyremote_bootstrap_finalize() -> PyremotePayloadRuntime:
822
830
  with open(os.environ.pop(_PYREMOTE_BOOTSTRAP_SRC_FILE_VAR)) as sf:
823
831
  main_src = sf.read()
824
832
 
825
- # Restore original argv0
833
+ # Restore vars
826
834
  sys.executable = os.environ.pop(_PYREMOTE_BOOTSTRAP_ARGV0_VAR)
827
-
828
- # Grab context name
829
835
  context_name = os.environ.pop(_PYREMOTE_BOOTSTRAP_CONTEXT_NAME_VAR)
830
836
 
831
837
  # Write third ack
@@ -1000,12 +1006,102 @@ class PyremoteBootstrapDriver:
1000
1006
  else:
1001
1007
  raise TypeError(go)
1002
1008
 
1009
+ async def async_run(
1010
+ self,
1011
+ input: ta.Any, # asyncio.StreamWriter # noqa
1012
+ output: ta.Any, # asyncio.StreamReader
1013
+ ) -> Result:
1014
+ gen = self.gen()
1015
+
1016
+ gi: ta.Optional[bytes] = None
1017
+ while True:
1018
+ try:
1019
+ if gi is not None:
1020
+ go = gen.send(gi)
1021
+ else:
1022
+ go = next(gen)
1023
+ except StopIteration as e:
1024
+ return e.value
1025
+
1026
+ if isinstance(go, self.Read):
1027
+ if len(gi := await input.read(go.sz)) != go.sz:
1028
+ raise EOFError
1029
+ elif isinstance(go, self.Write):
1030
+ gi = None
1031
+ output.write(go.d)
1032
+ await output.drain()
1033
+ else:
1034
+ raise TypeError(go)
1035
+
1036
+
1037
+ ########################################
1038
+ # ../../../omlish/lite/asyncio/asyncio.py
1039
+
1040
+
1041
+ ##
1042
+
1043
+
1044
+ ASYNCIO_DEFAULT_BUFFER_LIMIT = 2 ** 16
1045
+
1046
+
1047
+ async def asyncio_open_stream_reader(
1048
+ f: ta.IO,
1049
+ loop: ta.Any = None,
1050
+ *,
1051
+ limit: int = ASYNCIO_DEFAULT_BUFFER_LIMIT,
1052
+ ) -> asyncio.StreamReader:
1053
+ if loop is None:
1054
+ loop = asyncio.get_running_loop()
1055
+
1056
+ reader = asyncio.StreamReader(limit=limit, loop=loop)
1057
+ await loop.connect_read_pipe(
1058
+ lambda: asyncio.StreamReaderProtocol(reader, loop=loop),
1059
+ f,
1060
+ )
1061
+
1062
+ return reader
1063
+
1064
+
1065
+ async def asyncio_open_stream_writer(
1066
+ f: ta.IO,
1067
+ loop: ta.Any = None,
1068
+ ) -> asyncio.StreamWriter:
1069
+ if loop is None:
1070
+ loop = asyncio.get_running_loop()
1071
+
1072
+ writer_transport, writer_protocol = await loop.connect_write_pipe(
1073
+ lambda: asyncio.streams.FlowControlMixin(loop=loop),
1074
+ f,
1075
+ )
1076
+
1077
+ return asyncio.streams.StreamWriter(
1078
+ writer_transport,
1079
+ writer_protocol,
1080
+ None,
1081
+ loop,
1082
+ )
1083
+
1084
+
1085
+ ##
1086
+
1087
+
1088
+ def asyncio_maybe_timeout(
1089
+ fut: AwaitableT,
1090
+ timeout: ta.Optional[float] = None,
1091
+ ) -> AwaitableT:
1092
+ if timeout is not None:
1093
+ fut = asyncio.wait_for(fut, timeout) # type: ignore
1094
+ return fut
1095
+
1003
1096
 
1004
1097
  ########################################
1005
1098
  # ../../../omlish/lite/cached.py
1006
1099
 
1007
1100
 
1008
- class _cached_nullary: # noqa
1101
+ ##
1102
+
1103
+
1104
+ class _AbstractCachedNullary:
1009
1105
  def __init__(self, fn):
1010
1106
  super().__init__()
1011
1107
  self._fn = fn
@@ -1013,17 +1109,25 @@ class _cached_nullary: # noqa
1013
1109
  functools.update_wrapper(self, fn)
1014
1110
 
1015
1111
  def __call__(self, *args, **kwargs): # noqa
1016
- if self._value is self._missing:
1017
- self._value = self._fn()
1018
- return self._value
1112
+ raise TypeError
1019
1113
 
1020
1114
  def __get__(self, instance, owner): # noqa
1021
1115
  bound = instance.__dict__[self._fn.__name__] = self.__class__(self._fn.__get__(instance, owner))
1022
1116
  return bound
1023
1117
 
1024
1118
 
1119
+ ##
1120
+
1121
+
1122
+ class _CachedNullary(_AbstractCachedNullary):
1123
+ def __call__(self, *args, **kwargs): # noqa
1124
+ if self._value is self._missing:
1125
+ self._value = self._fn()
1126
+ return self._value
1127
+
1128
+
1025
1129
  def cached_nullary(fn): # ta.Callable[..., T]) -> ta.Callable[..., T]:
1026
- return _cached_nullary(fn)
1130
+ return _CachedNullary(fn)
1027
1131
 
1028
1132
 
1029
1133
  def static_init(fn: CallableT) -> CallableT:
@@ -1032,6 +1136,20 @@ def static_init(fn: CallableT) -> CallableT:
1032
1136
  return fn
1033
1137
 
1034
1138
 
1139
+ ##
1140
+
1141
+
1142
+ class _AsyncCachedNullary(_AbstractCachedNullary):
1143
+ async def __call__(self, *args, **kwargs):
1144
+ if self._value is self._missing:
1145
+ self._value = await self._fn()
1146
+ return self._value
1147
+
1148
+
1149
+ def async_cached_nullary(fn): # ta.Callable[..., T]) -> ta.Callable[..., T]:
1150
+ return _AsyncCachedNullary(fn)
1151
+
1152
+
1035
1153
  ########################################
1036
1154
  # ../../../omlish/lite/check.py
1037
1155
 
@@ -1129,6 +1247,30 @@ def check_non_empty(v: SizedT) -> SizedT:
1129
1247
  return v
1130
1248
 
1131
1249
 
1250
+ ########################################
1251
+ # ../../../omlish/lite/deathsig.py
1252
+
1253
+
1254
+ LINUX_PR_SET_PDEATHSIG = 1 # Second arg is a signal
1255
+ LINUX_PR_GET_PDEATHSIG = 2 # Second arg is a ptr to return the signal
1256
+
1257
+
1258
+ def set_process_deathsig(sig: int) -> bool:
1259
+ if sys.platform == 'linux':
1260
+ libc = ct.CDLL('libc.so.6')
1261
+
1262
+ # int prctl(int option, unsigned long arg2, unsigned long arg3, unsigned long arg4, unsigned long arg5);
1263
+ libc.prctl.restype = ct.c_int
1264
+ libc.prctl.argtypes = [ct.c_int, ct.c_ulong, ct.c_ulong, ct.c_ulong, ct.c_ulong]
1265
+
1266
+ libc.prctl(LINUX_PR_SET_PDEATHSIG, sig, 0, 0, 0, 0)
1267
+
1268
+ return True
1269
+
1270
+ else:
1271
+ return False
1272
+
1273
+
1132
1274
  ########################################
1133
1275
  # ../../../omlish/lite/json.py
1134
1276
 
@@ -1910,8 +2052,8 @@ class Command(abc.ABC, ta.Generic[CommandOutputT]):
1910
2052
  pass
1911
2053
 
1912
2054
  @ta.final
1913
- def execute(self, executor: 'CommandExecutor') -> CommandOutputT:
1914
- return check_isinstance(executor.execute(self), self.Output) # type: ignore[return-value]
2055
+ async def execute(self, executor: 'CommandExecutor') -> CommandOutputT:
2056
+ return check_isinstance(await executor.execute(self), self.Output) # type: ignore[return-value]
1915
2057
 
1916
2058
 
1917
2059
  ##
@@ -1972,10 +2114,10 @@ class CommandOutputOrExceptionData(CommandOutputOrException):
1972
2114
 
1973
2115
  class CommandExecutor(abc.ABC, ta.Generic[CommandT, CommandOutputT]):
1974
2116
  @abc.abstractmethod
1975
- def execute(self, cmd: CommandT) -> CommandOutputT:
2117
+ def execute(self, cmd: CommandT) -> ta.Awaitable[CommandOutputT]:
1976
2118
  raise NotImplementedError
1977
2119
 
1978
- def try_execute(
2120
+ async def try_execute(
1979
2121
  self,
1980
2122
  cmd: CommandT,
1981
2123
  *,
@@ -1983,7 +2125,7 @@ class CommandExecutor(abc.ABC, ta.Generic[CommandT, CommandOutputT]):
1983
2125
  omit_exc_object: bool = False,
1984
2126
  ) -> CommandOutputOrException[CommandOutputT]:
1985
2127
  try:
1986
- o = self.execute(cmd)
2128
+ o = await self.execute(cmd)
1987
2129
 
1988
2130
  except Exception as e: # noqa
1989
2131
  if log is not None:
@@ -2054,8 +2196,18 @@ def build_command_name_map(crs: CommandRegistrations) -> CommandNameMap:
2054
2196
  class RemoteConfig:
2055
2197
  payload_file: ta.Optional[str] = None
2056
2198
 
2199
+ set_pgid: bool = True
2200
+
2201
+ deathsig: ta.Optional[str] = 'KILL'
2202
+
2057
2203
  pycharm_remote_debug: ta.Optional[PycharmRemoteDebug] = None
2058
2204
 
2205
+ forward_logging: bool = True
2206
+
2207
+ timebomb_delay_s: ta.Optional[float] = 60 * 60.
2208
+
2209
+ heartbeat_interval_s: float = 3.
2210
+
2059
2211
 
2060
2212
  ########################################
2061
2213
  # ../remote/payload.py
@@ -3788,6 +3940,8 @@ class InterpSpecifier:
3788
3940
  s, o = InterpOpts.parse_suffix(s)
3789
3941
  if not any(s.startswith(o) for o in Specifier.OPERATORS):
3790
3942
  s = '~=' + s
3943
+ if s.count('.') < 2:
3944
+ s += '.0'
3791
3945
  return cls(
3792
3946
  specifier=Specifier(s),
3793
3947
  opts=o,
@@ -3834,9 +3988,9 @@ class LocalCommandExecutor(CommandExecutor):
3834
3988
 
3835
3989
  self._command_executors = command_executors
3836
3990
 
3837
- def execute(self, cmd: Command) -> Command.Output:
3991
+ async def execute(self, cmd: Command) -> Command.Output:
3838
3992
  ce: CommandExecutor = self._command_executors[type(cmd)]
3839
- return ce.execute(cmd)
3993
+ return await ce.execute(cmd)
3840
3994
 
3841
3995
 
3842
3996
  ########################################
@@ -3882,7 +4036,7 @@ class DeployCommand(Command['DeployCommand.Output']):
3882
4036
 
3883
4037
 
3884
4038
  class DeployCommandExecutor(CommandExecutor[DeployCommand, DeployCommand.Output]):
3885
- def execute(self, cmd: DeployCommand) -> DeployCommand.Output:
4039
+ async def execute(self, cmd: DeployCommand) -> DeployCommand.Output:
3886
4040
  log.info('Deploying!')
3887
4041
 
3888
4042
  return DeployCommand.Output()
@@ -3904,11 +4058,30 @@ ObjMarshalerInstallers = ta.NewType('ObjMarshalerInstallers', ta.Sequence[ObjMar
3904
4058
  # ../remote/channel.py
3905
4059
 
3906
4060
 
3907
- class RemoteChannel:
4061
+ ##
4062
+
4063
+
4064
+ class RemoteChannel(abc.ABC):
4065
+ @abc.abstractmethod
4066
+ def send_obj(self, o: ta.Any, ty: ta.Any = None) -> ta.Awaitable[None]:
4067
+ raise NotImplementedError
4068
+
4069
+ @abc.abstractmethod
4070
+ def recv_obj(self, ty: ta.Type[T]) -> ta.Awaitable[ta.Optional[T]]:
4071
+ raise NotImplementedError
4072
+
4073
+ def set_marshaler(self, msh: ObjMarshalerManager) -> None: # noqa
4074
+ pass
4075
+
4076
+
4077
+ ##
4078
+
4079
+
4080
+ class RemoteChannelImpl(RemoteChannel):
3908
4081
  def __init__(
3909
4082
  self,
3910
- input: ta.IO, # noqa
3911
- output: ta.IO,
4083
+ input: asyncio.StreamReader, # noqa
4084
+ output: asyncio.StreamWriter,
3912
4085
  *,
3913
4086
  msh: ObjMarshalerManager = OBJ_MARSHALER_MANAGER,
3914
4087
  ) -> None:
@@ -3918,41 +4091,46 @@ class RemoteChannel:
3918
4091
  self._output = output
3919
4092
  self._msh = msh
3920
4093
 
3921
- self._lock = threading.RLock()
4094
+ self._input_lock = asyncio.Lock()
4095
+ self._output_lock = asyncio.Lock()
3922
4096
 
3923
4097
  def set_marshaler(self, msh: ObjMarshalerManager) -> None:
3924
4098
  self._msh = msh
3925
4099
 
3926
- def _send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
4100
+ #
4101
+
4102
+ async def _send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
3927
4103
  j = json_dumps_compact(self._msh.marshal_obj(o, ty))
3928
4104
  d = j.encode('utf-8')
3929
4105
 
3930
4106
  self._output.write(struct.pack('<I', len(d)))
3931
4107
  self._output.write(d)
3932
- self._output.flush()
4108
+ await self._output.drain()
3933
4109
 
3934
- def send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
3935
- with self._lock:
3936
- return self._send_obj(o, ty)
4110
+ async def send_obj(self, o: ta.Any, ty: ta.Any = None) -> None:
4111
+ async with self._output_lock:
4112
+ return await self._send_obj(o, ty)
4113
+
4114
+ #
3937
4115
 
3938
- def _recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
3939
- d = self._input.read(4)
4116
+ async def _recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
4117
+ d = await self._input.read(4)
3940
4118
  if not d:
3941
4119
  return None
3942
4120
  if len(d) != 4:
3943
4121
  raise EOFError
3944
4122
 
3945
4123
  sz = struct.unpack('<I', d)[0]
3946
- d = self._input.read(sz)
4124
+ d = await self._input.read(sz)
3947
4125
  if len(d) != sz:
3948
4126
  raise EOFError
3949
4127
 
3950
4128
  j = json.loads(d.decode('utf-8'))
3951
4129
  return self._msh.unmarshal_obj(j, ty)
3952
4130
 
3953
- def recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
3954
- with self._lock:
3955
- return self._recv_obj(ty)
4131
+ async def recv_obj(self, ty: ta.Type[T]) -> ta.Optional[T]:
4132
+ async with self._input_lock:
4133
+ return await self._recv_obj(ty)
3956
4134
 
3957
4135
 
3958
4136
  ########################################
@@ -3986,7 +4164,7 @@ def subprocess_maybe_shell_wrap_exec(*args: str) -> ta.Tuple[str, ...]:
3986
4164
  return args
3987
4165
 
3988
4166
 
3989
- def _prepare_subprocess_invocation(
4167
+ def prepare_subprocess_invocation(
3990
4168
  *args: str,
3991
4169
  env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
3992
4170
  extra_env: ta.Optional[ta.Mapping[str, ta.Any]] = None,
@@ -3994,9 +4172,9 @@ def _prepare_subprocess_invocation(
3994
4172
  shell: bool = False,
3995
4173
  **kwargs: ta.Any,
3996
4174
  ) -> ta.Tuple[ta.Tuple[ta.Any, ...], ta.Dict[str, ta.Any]]:
3997
- log.debug(args)
4175
+ log.debug('prepare_subprocess_invocation: args=%r', args)
3998
4176
  if extra_env:
3999
- log.debug(extra_env)
4177
+ log.debug('prepare_subprocess_invocation: extra_env=%r', extra_env)
4000
4178
 
4001
4179
  if extra_env:
4002
4180
  env = {**(env if env is not None else os.environ), **extra_env}
@@ -4015,14 +4193,46 @@ def _prepare_subprocess_invocation(
4015
4193
  )
4016
4194
 
4017
4195
 
4018
- def subprocess_check_call(*args: str, stdout=sys.stderr, **kwargs: ta.Any) -> None:
4019
- args, kwargs = _prepare_subprocess_invocation(*args, stdout=stdout, **kwargs)
4020
- return subprocess.check_call(args, **kwargs) # type: ignore
4196
+ ##
4197
+
4198
+
4199
+ @contextlib.contextmanager
4200
+ def subprocess_common_context(*args: ta.Any, **kwargs: ta.Any) -> ta.Iterator[None]:
4201
+ start_time = time.time()
4202
+ try:
4203
+ log.debug('subprocess_common_context.try: args=%r', args)
4204
+ yield
4205
+
4206
+ except Exception as exc: # noqa
4207
+ log.debug('subprocess_common_context.except: exc=%r', exc)
4208
+ raise
4209
+
4210
+ finally:
4211
+ end_time = time.time()
4212
+ elapsed_s = end_time - start_time
4213
+ log.debug('subprocess_common_context.finally: elapsed_s=%f args=%r', elapsed_s, args)
4214
+
4215
+
4216
+ ##
4217
+
4218
+
4219
+ def subprocess_check_call(
4220
+ *args: str,
4221
+ stdout: ta.Any = sys.stderr,
4222
+ **kwargs: ta.Any,
4223
+ ) -> None:
4224
+ args, kwargs = prepare_subprocess_invocation(*args, stdout=stdout, **kwargs)
4225
+ with subprocess_common_context(*args, **kwargs):
4226
+ return subprocess.check_call(args, **kwargs) # type: ignore
4021
4227
 
4022
4228
 
4023
- def subprocess_check_output(*args: str, **kwargs: ta.Any) -> bytes:
4024
- args, kwargs = _prepare_subprocess_invocation(*args, **kwargs)
4025
- return subprocess.check_output(args, **kwargs)
4229
+ def subprocess_check_output(
4230
+ *args: str,
4231
+ **kwargs: ta.Any,
4232
+ ) -> bytes:
4233
+ args, kwargs = prepare_subprocess_invocation(*args, **kwargs)
4234
+ with subprocess_common_context(*args, **kwargs):
4235
+ return subprocess.check_output(args, **kwargs)
4026
4236
 
4027
4237
 
4028
4238
  def subprocess_check_output_str(*args: str, **kwargs: ta.Any) -> str:
@@ -4038,16 +4248,31 @@ DEFAULT_SUBPROCESS_TRY_EXCEPTIONS: ta.Tuple[ta.Type[Exception], ...] = (
4038
4248
  )
4039
4249
 
4040
4250
 
4041
- def subprocess_try_call(
4042
- *args: str,
4251
+ def _subprocess_try_run(
4252
+ fn: ta.Callable[..., T],
4253
+ *args: ta.Any,
4043
4254
  try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
4044
4255
  **kwargs: ta.Any,
4045
- ) -> bool:
4256
+ ) -> ta.Union[T, Exception]:
4046
4257
  try:
4047
- subprocess_check_call(*args, **kwargs)
4258
+ return fn(*args, **kwargs)
4048
4259
  except try_exceptions as e: # noqa
4049
4260
  if log.isEnabledFor(logging.DEBUG):
4050
4261
  log.exception('command failed')
4262
+ return e
4263
+
4264
+
4265
+ def subprocess_try_call(
4266
+ *args: str,
4267
+ try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
4268
+ **kwargs: ta.Any,
4269
+ ) -> bool:
4270
+ if isinstance(_subprocess_try_run(
4271
+ subprocess_check_call,
4272
+ *args,
4273
+ try_exceptions=try_exceptions,
4274
+ **kwargs,
4275
+ ), Exception):
4051
4276
  return False
4052
4277
  else:
4053
4278
  return True
@@ -4058,12 +4283,15 @@ def subprocess_try_output(
4058
4283
  try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
4059
4284
  **kwargs: ta.Any,
4060
4285
  ) -> ta.Optional[bytes]:
4061
- try:
4062
- return subprocess_check_output(*args, **kwargs)
4063
- except try_exceptions as e: # noqa
4064
- if log.isEnabledFor(logging.DEBUG):
4065
- log.exception('command failed')
4286
+ if isinstance(ret := _subprocess_try_run(
4287
+ subprocess_check_output,
4288
+ *args,
4289
+ try_exceptions=try_exceptions,
4290
+ **kwargs,
4291
+ ), Exception):
4066
4292
  return None
4293
+ else:
4294
+ return ret
4067
4295
 
4068
4296
 
4069
4297
  def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
@@ -4090,175 +4318,922 @@ def subprocess_close(
4090
4318
 
4091
4319
 
4092
4320
  ########################################
4093
- # ../../../omdev/interp/inspect.py
4321
+ # ../remote/execution.py
4094
4322
 
4095
4323
 
4096
- @dc.dataclass(frozen=True)
4097
- class InterpInspection:
4098
- exe: str
4099
- version: Version
4324
+ ##
4100
4325
 
4101
- version_str: str
4102
- config_vars: ta.Mapping[str, str]
4103
- prefix: str
4104
- base_prefix: str
4105
4326
 
4106
- @property
4107
- def opts(self) -> InterpOpts:
4108
- return InterpOpts(
4109
- threaded=bool(self.config_vars.get('Py_GIL_DISABLED')),
4110
- debug=bool(self.config_vars.get('Py_DEBUG')),
4111
- )
4327
+ class _RemoteProtocol:
4328
+ class Message(abc.ABC): # noqa
4329
+ async def send(self, chan: RemoteChannel) -> None:
4330
+ await chan.send_obj(self, _RemoteProtocol.Message)
4112
4331
 
4113
- @property
4114
- def iv(self) -> InterpVersion:
4115
- return InterpVersion(
4116
- version=self.version,
4117
- opts=self.opts,
4118
- )
4332
+ @classmethod
4333
+ async def recv(cls: ta.Type[T], chan: RemoteChannel) -> ta.Optional[T]:
4334
+ return await chan.recv_obj(cls)
4119
4335
 
4120
- @property
4121
- def is_venv(self) -> bool:
4122
- return self.prefix != self.base_prefix
4336
+ #
4123
4337
 
4338
+ class Request(Message, abc.ABC): # noqa
4339
+ pass
4124
4340
 
4125
- class InterpInspector:
4341
+ @dc.dataclass(frozen=True)
4342
+ class CommandRequest(Request):
4343
+ seq: int
4344
+ cmd: Command
4126
4345
 
4127
- def __init__(self) -> None:
4128
- super().__init__()
4346
+ @dc.dataclass(frozen=True)
4347
+ class PingRequest(Request):
4348
+ time: float
4129
4349
 
4130
- self._cache: ta.Dict[str, ta.Optional[InterpInspection]] = {}
4350
+ #
4131
4351
 
4132
- _RAW_INSPECTION_CODE = """
4133
- __import__('json').dumps(dict(
4134
- version_str=__import__('sys').version,
4135
- prefix=__import__('sys').prefix,
4136
- base_prefix=__import__('sys').base_prefix,
4137
- config_vars=__import__('sysconfig').get_config_vars(),
4138
- ))"""
4352
+ class Response(Message, abc.ABC): # noqa
4353
+ pass
4139
4354
 
4140
- _INSPECTION_CODE = ''.join(l.strip() for l in _RAW_INSPECTION_CODE.splitlines())
4355
+ @dc.dataclass(frozen=True)
4356
+ class LogResponse(Response):
4357
+ s: str
4141
4358
 
4142
- @staticmethod
4143
- def _build_inspection(
4144
- exe: str,
4145
- output: str,
4146
- ) -> InterpInspection:
4147
- dct = json.loads(output)
4359
+ @dc.dataclass(frozen=True)
4360
+ class CommandResponse(Response):
4361
+ seq: int
4362
+ res: CommandOutputOrExceptionData
4148
4363
 
4149
- version = Version(dct['version_str'].split()[0])
4364
+ @dc.dataclass(frozen=True)
4365
+ class PingResponse(Response):
4366
+ time: float
4150
4367
 
4151
- return InterpInspection(
4152
- exe=exe,
4153
- version=version,
4154
- **{k: dct[k] for k in (
4155
- 'version_str',
4156
- 'prefix',
4157
- 'base_prefix',
4158
- 'config_vars',
4159
- )},
4160
- )
4161
4368
 
4162
- @classmethod
4163
- def running(cls) -> 'InterpInspection':
4164
- return cls._build_inspection(sys.executable, eval(cls._INSPECTION_CODE)) # noqa
4369
+ ##
4165
4370
 
4166
- def _inspect(self, exe: str) -> InterpInspection:
4167
- output = subprocess_check_output(exe, '-c', f'print({self._INSPECTION_CODE})', quiet=True)
4168
- return self._build_inspection(exe, output.decode())
4169
4371
 
4170
- def inspect(self, exe: str) -> ta.Optional[InterpInspection]:
4171
- try:
4172
- return self._cache[exe]
4173
- except KeyError:
4174
- ret: ta.Optional[InterpInspection]
4175
- try:
4176
- ret = self._inspect(exe)
4177
- except Exception as e: # noqa
4178
- if log.isEnabledFor(logging.DEBUG):
4179
- log.exception('Failed to inspect interp: %s', exe)
4180
- ret = None
4181
- self._cache[exe] = ret
4182
- return ret
4372
+ class _RemoteLogHandler(logging.Handler):
4373
+ def __init__(
4374
+ self,
4375
+ chan: RemoteChannel,
4376
+ loop: ta.Any = None,
4377
+ ) -> None:
4378
+ super().__init__()
4183
4379
 
4380
+ self._chan = chan
4381
+ self._loop = loop
4184
4382
 
4185
- INTERP_INSPECTOR = InterpInspector()
4383
+ def emit(self, record):
4384
+ msg = self.format(record)
4186
4385
 
4386
+ async def inner():
4387
+ await _RemoteProtocol.LogResponse(msg).send(self._chan)
4187
4388
 
4188
- ########################################
4189
- # ../commands/subprocess.py
4389
+ loop = self._loop
4390
+ if loop is None:
4391
+ loop = asyncio.get_running_loop()
4392
+ if loop is not None:
4393
+ asyncio.run_coroutine_threadsafe(inner(), loop)
4190
4394
 
4191
4395
 
4192
4396
  ##
4193
4397
 
4194
4398
 
4195
- @dc.dataclass(frozen=True)
4196
- class SubprocessCommand(Command['SubprocessCommand.Output']):
4197
- cmd: ta.Sequence[str]
4399
+ class _RemoteCommandHandler:
4400
+ def __init__(
4401
+ self,
4402
+ chan: RemoteChannel,
4403
+ executor: CommandExecutor,
4404
+ *,
4405
+ stop: ta.Optional[asyncio.Event] = None,
4406
+ ) -> None:
4407
+ super().__init__()
4198
4408
 
4199
- shell: bool = False
4200
- cwd: ta.Optional[str] = None
4201
- env: ta.Optional[ta.Mapping[str, str]] = None
4409
+ self._chan = chan
4410
+ self._executor = executor
4411
+ self._stop = stop if stop is not None else asyncio.Event()
4202
4412
 
4203
- stdout: str = 'pipe' # SubprocessChannelOption
4204
- stderr: str = 'pipe' # SubprocessChannelOption
4413
+ self._cmds_by_seq: ta.Dict[int, _RemoteCommandHandler._Command] = {}
4205
4414
 
4206
- input: ta.Optional[bytes] = None
4207
- timeout: ta.Optional[float] = None
4415
+ @dc.dataclass(frozen=True)
4416
+ class _Command:
4417
+ req: _RemoteProtocol.CommandRequest
4418
+ fut: asyncio.Future
4419
+
4420
+ async def run(self) -> None:
4421
+ stop_task = asyncio.create_task(self._stop.wait())
4422
+ recv_task: ta.Optional[asyncio.Task] = None
4423
+
4424
+ while not self._stop.is_set():
4425
+ if recv_task is None:
4426
+ recv_task = asyncio.create_task(_RemoteProtocol.Request.recv(self._chan))
4427
+
4428
+ done, pending = await asyncio.wait([
4429
+ stop_task,
4430
+ recv_task,
4431
+ ], return_when=asyncio.FIRST_COMPLETED)
4432
+
4433
+ if recv_task in done:
4434
+ msg: ta.Optional[_RemoteProtocol.Message] = check_isinstance(
4435
+ recv_task.result(),
4436
+ (_RemoteProtocol.Message, type(None)),
4437
+ )
4438
+ recv_task = None
4439
+
4440
+ if msg is None:
4441
+ break
4442
+
4443
+ await self._handle_message(msg)
4444
+
4445
+ async def _handle_message(self, msg: _RemoteProtocol.Message) -> None:
4446
+ if isinstance(msg, _RemoteProtocol.PingRequest):
4447
+ log.debug('Ping: %r', msg)
4448
+ await _RemoteProtocol.PingResponse(
4449
+ time=msg.time,
4450
+ ).send(self._chan)
4451
+
4452
+ elif isinstance(msg, _RemoteProtocol.CommandRequest):
4453
+ fut = asyncio.create_task(self._handle_command_request(msg))
4454
+ self._cmds_by_seq[msg.seq] = _RemoteCommandHandler._Command(
4455
+ req=msg,
4456
+ fut=fut,
4457
+ )
4458
+
4459
+ else:
4460
+ raise TypeError(msg)
4461
+
4462
+ async def _handle_command_request(self, req: _RemoteProtocol.CommandRequest) -> None:
4463
+ res = await self._executor.try_execute(
4464
+ req.cmd,
4465
+ log=log,
4466
+ omit_exc_object=True,
4467
+ )
4468
+
4469
+ await _RemoteProtocol.CommandResponse(
4470
+ seq=req.seq,
4471
+ res=CommandOutputOrExceptionData(
4472
+ output=res.output,
4473
+ exception=res.exception,
4474
+ ),
4475
+ ).send(self._chan)
4476
+
4477
+ self._cmds_by_seq.pop(req.seq) # noqa
4478
+
4479
+
4480
+ ##
4481
+
4482
+
4483
+ @dc.dataclass()
4484
+ class RemoteCommandError(Exception):
4485
+ e: CommandException
4486
+
4487
+
4488
+ class RemoteCommandExecutor(CommandExecutor):
4489
+ def __init__(self, chan: RemoteChannel) -> None:
4490
+ super().__init__()
4491
+
4492
+ self._chan = chan
4493
+
4494
+ self._cmd_seq = itertools.count()
4495
+ self._queue: asyncio.Queue = asyncio.Queue() # asyncio.Queue[RemoteCommandExecutor._Request]
4496
+ self._stop = asyncio.Event()
4497
+ self._loop_task: ta.Optional[asyncio.Task] = None
4498
+ self._reqs_by_seq: ta.Dict[int, RemoteCommandExecutor._Request] = {}
4499
+
4500
+ #
4501
+
4502
+ async def start(self) -> None:
4503
+ check_none(self._loop_task)
4504
+ check_state(not self._stop.is_set())
4505
+ self._loop_task = asyncio.create_task(self._loop())
4506
+
4507
+ async def aclose(self) -> None:
4508
+ self._stop.set()
4509
+ if self._loop_task is not None:
4510
+ await self._loop_task
4511
+
4512
+ #
4513
+
4514
+ @dc.dataclass(frozen=True)
4515
+ class _Request:
4516
+ seq: int
4517
+ cmd: Command
4518
+ fut: asyncio.Future
4519
+
4520
+ async def _loop(self) -> None:
4521
+ log.debug('RemoteCommandExecutor loop start: %r', self)
4522
+
4523
+ stop_task = asyncio.create_task(self._stop.wait())
4524
+ queue_task: ta.Optional[asyncio.Task] = None
4525
+ recv_task: ta.Optional[asyncio.Task] = None
4526
+
4527
+ while not self._stop.is_set():
4528
+ if queue_task is None:
4529
+ queue_task = asyncio.create_task(self._queue.get())
4530
+ if recv_task is None:
4531
+ recv_task = asyncio.create_task(_RemoteProtocol.Message.recv(self._chan))
4532
+
4533
+ done, pending = await asyncio.wait([
4534
+ stop_task,
4535
+ queue_task,
4536
+ recv_task,
4537
+ ], return_when=asyncio.FIRST_COMPLETED)
4538
+
4539
+ if queue_task in done:
4540
+ req = check_isinstance(queue_task.result(), RemoteCommandExecutor._Request)
4541
+ queue_task = None
4542
+ await self._handle_request(req)
4543
+
4544
+ if recv_task in done:
4545
+ msg: ta.Optional[_RemoteProtocol.Message] = check_isinstance(
4546
+ recv_task.result(),
4547
+ (_RemoteProtocol.Message, type(None)),
4548
+ )
4549
+ recv_task = None
4550
+
4551
+ if msg is None:
4552
+ log.debug('RemoteCommandExecutor got eof: %r', self)
4553
+ break
4554
+
4555
+ await self._handle_message(msg)
4556
+
4557
+ log.debug('RemoteCommandExecutor loop stopping: %r', self)
4558
+
4559
+ for task in [
4560
+ stop_task,
4561
+ queue_task,
4562
+ recv_task,
4563
+ ]:
4564
+ if task is not None and not task.done():
4565
+ task.cancel()
4566
+
4567
+ for req in self._reqs_by_seq.values():
4568
+ req.fut.cancel()
4569
+
4570
+ log.debug('RemoteCommandExecutor loop exited: %r', self)
4571
+
4572
+ async def _handle_request(self, req: _Request) -> None:
4573
+ self._reqs_by_seq[req.seq] = req
4574
+ await _RemoteProtocol.CommandRequest(
4575
+ seq=req.seq,
4576
+ cmd=req.cmd,
4577
+ ).send(self._chan)
4578
+
4579
+ async def _handle_message(self, msg: _RemoteProtocol.Message) -> None:
4580
+ if isinstance(msg, _RemoteProtocol.PingRequest):
4581
+ log.debug('Ping: %r', msg)
4582
+ await _RemoteProtocol.PingResponse(
4583
+ time=msg.time,
4584
+ ).send(self._chan)
4585
+
4586
+ elif isinstance(msg, _RemoteProtocol.LogResponse):
4587
+ log.info(msg.s)
4588
+
4589
+ elif isinstance(msg, _RemoteProtocol.CommandResponse):
4590
+ req = self._reqs_by_seq.pop(msg.seq)
4591
+ req.fut.set_result(msg.res)
4592
+
4593
+ else:
4594
+ raise TypeError(msg)
4595
+
4596
+ #
4597
+
4598
+ async def _remote_execute(self, cmd: Command) -> CommandOutputOrException:
4599
+ req = RemoteCommandExecutor._Request(
4600
+ seq=next(self._cmd_seq),
4601
+ cmd=cmd,
4602
+ fut=asyncio.Future(),
4603
+ )
4604
+ await self._queue.put(req)
4605
+ return await req.fut
4606
+
4607
+ # @ta.override
4608
+ async def execute(self, cmd: Command) -> Command.Output:
4609
+ r = await self._remote_execute(cmd)
4610
+ if (e := r.exception) is not None:
4611
+ raise RemoteCommandError(e)
4612
+ else:
4613
+ return check_not_none(r.output)
4614
+
4615
+ # @ta.override
4616
+ async def try_execute(
4617
+ self,
4618
+ cmd: Command,
4619
+ *,
4620
+ log: ta.Optional[logging.Logger] = None,
4621
+ omit_exc_object: bool = False,
4622
+ ) -> CommandOutputOrException:
4623
+ try:
4624
+ r = await self._remote_execute(cmd)
4625
+
4626
+ except Exception as e: # noqa
4627
+ if log is not None:
4628
+ log.exception('Exception executing remote command: %r', type(cmd))
4629
+
4630
+ return CommandOutputOrExceptionData(exception=CommandException.of(
4631
+ e,
4632
+ omit_exc_object=omit_exc_object,
4633
+ cmd=cmd,
4634
+ ))
4635
+
4636
+ else:
4637
+ return r
4638
+
4639
+
4640
+ ########################################
4641
+ # ../../../omlish/lite/asyncio/subprocesses.py
4642
+
4643
+
4644
+ ##
4645
+
4646
+
4647
+ @contextlib.asynccontextmanager
4648
+ async def asyncio_subprocess_popen(
4649
+ *cmd: str,
4650
+ shell: bool = False,
4651
+ timeout: ta.Optional[float] = None,
4652
+ **kwargs: ta.Any,
4653
+ ) -> ta.AsyncGenerator[asyncio.subprocess.Process, None]:
4654
+ fac: ta.Any
4655
+ if shell:
4656
+ fac = functools.partial(
4657
+ asyncio.create_subprocess_shell,
4658
+ check_single(cmd),
4659
+ )
4660
+ else:
4661
+ fac = functools.partial(
4662
+ asyncio.create_subprocess_exec,
4663
+ *cmd,
4664
+ )
4665
+
4666
+ with subprocess_common_context(
4667
+ *cmd,
4668
+ shell=shell,
4669
+ timeout=timeout,
4670
+ **kwargs,
4671
+ ):
4672
+ proc: asyncio.subprocess.Process
4673
+ proc = await fac(**kwargs)
4674
+ try:
4675
+ yield proc
4676
+
4677
+ finally:
4678
+ await asyncio_maybe_timeout(proc.wait(), timeout)
4679
+
4680
+
4681
+ ##
4682
+
4683
+
4684
+ class AsyncioProcessCommunicator:
4685
+ def __init__(
4686
+ self,
4687
+ proc: asyncio.subprocess.Process,
4688
+ loop: ta.Optional[ta.Any] = None,
4689
+ ) -> None:
4690
+ super().__init__()
4691
+
4692
+ if loop is None:
4693
+ loop = asyncio.get_running_loop()
4694
+
4695
+ self._proc = proc
4696
+ self._loop = loop
4697
+
4698
+ self._transport: asyncio.base_subprocess.BaseSubprocessTransport = check_isinstance(
4699
+ proc._transport, # type: ignore # noqa
4700
+ asyncio.base_subprocess.BaseSubprocessTransport,
4701
+ )
4702
+
4703
+ @property
4704
+ def _debug(self) -> bool:
4705
+ return self._loop.get_debug()
4706
+
4707
+ async def _feed_stdin(self, input: bytes) -> None: # noqa
4708
+ stdin = check_not_none(self._proc.stdin)
4709
+ try:
4710
+ if input is not None:
4711
+ stdin.write(input)
4712
+ if self._debug:
4713
+ log.debug('%r communicate: feed stdin (%s bytes)', self, len(input))
4714
+
4715
+ await stdin.drain()
4716
+
4717
+ except (BrokenPipeError, ConnectionResetError) as exc:
4718
+ # communicate() ignores BrokenPipeError and ConnectionResetError. write() and drain() can raise these
4719
+ # exceptions.
4720
+ if self._debug:
4721
+ log.debug('%r communicate: stdin got %r', self, exc)
4722
+
4723
+ if self._debug:
4724
+ log.debug('%r communicate: close stdin', self)
4725
+
4726
+ stdin.close()
4727
+
4728
+ async def _noop(self) -> None:
4729
+ return None
4730
+
4731
+ async def _read_stream(self, fd: int) -> bytes:
4732
+ transport: ta.Any = check_not_none(self._transport.get_pipe_transport(fd))
4733
+
4734
+ if fd == 2:
4735
+ stream = check_not_none(self._proc.stderr)
4736
+ else:
4737
+ check_equal(fd, 1)
4738
+ stream = check_not_none(self._proc.stdout)
4739
+
4740
+ if self._debug:
4741
+ name = 'stdout' if fd == 1 else 'stderr'
4742
+ log.debug('%r communicate: read %s', self, name)
4743
+
4744
+ output = await stream.read()
4745
+
4746
+ if self._debug:
4747
+ name = 'stdout' if fd == 1 else 'stderr'
4748
+ log.debug('%r communicate: close %s', self, name)
4749
+
4750
+ transport.close()
4751
+
4752
+ return output
4753
+
4754
+ class Communication(ta.NamedTuple):
4755
+ stdout: ta.Optional[bytes]
4756
+ stderr: ta.Optional[bytes]
4757
+
4758
+ async def _communicate(
4759
+ self,
4760
+ input: ta.Any = None, # noqa
4761
+ ) -> Communication:
4762
+ stdin_fut: ta.Any
4763
+ if self._proc.stdin is not None:
4764
+ stdin_fut = self._feed_stdin(input)
4765
+ else:
4766
+ stdin_fut = self._noop()
4767
+
4768
+ stdout_fut: ta.Any
4769
+ if self._proc.stdout is not None:
4770
+ stdout_fut = self._read_stream(1)
4771
+ else:
4772
+ stdout_fut = self._noop()
4773
+
4774
+ stderr_fut: ta.Any
4775
+ if self._proc.stderr is not None:
4776
+ stderr_fut = self._read_stream(2)
4777
+ else:
4778
+ stderr_fut = self._noop()
4779
+
4780
+ stdin_res, stdout_res, stderr_res = await asyncio.gather(stdin_fut, stdout_fut, stderr_fut)
4781
+
4782
+ await self._proc.wait()
4783
+
4784
+ return AsyncioProcessCommunicator.Communication(stdout_res, stderr_res)
4785
+
4786
+ async def communicate(
4787
+ self,
4788
+ input: ta.Any = None, # noqa
4789
+ timeout: ta.Optional[float] = None,
4790
+ ) -> Communication:
4791
+ return await asyncio_maybe_timeout(self._communicate(input), timeout)
4792
+
4793
+
4794
+ async def asyncio_subprocess_communicate(
4795
+ proc: asyncio.subprocess.Process,
4796
+ input: ta.Any = None, # noqa
4797
+ timeout: ta.Optional[float] = None,
4798
+ ) -> ta.Tuple[ta.Optional[bytes], ta.Optional[bytes]]:
4799
+ return await AsyncioProcessCommunicator(proc).communicate(input, timeout) # noqa
4800
+
4801
+
4802
+ ##
4803
+
4804
+
4805
+ async def _asyncio_subprocess_check_run(
4806
+ *args: str,
4807
+ input: ta.Any = None, # noqa
4808
+ timeout: ta.Optional[float] = None,
4809
+ **kwargs: ta.Any,
4810
+ ) -> ta.Tuple[ta.Optional[bytes], ta.Optional[bytes]]:
4811
+ args, kwargs = prepare_subprocess_invocation(*args, **kwargs)
4812
+
4813
+ proc: asyncio.subprocess.Process
4814
+ async with asyncio_subprocess_popen(*args, **kwargs) as proc:
4815
+ stdout, stderr = await asyncio_subprocess_communicate(proc, input, timeout)
4816
+
4817
+ if proc.returncode:
4818
+ raise subprocess.CalledProcessError(
4819
+ proc.returncode,
4820
+ args,
4821
+ output=stdout,
4822
+ stderr=stderr,
4823
+ )
4824
+
4825
+ return stdout, stderr
4826
+
4827
+
4828
+ async def asyncio_subprocess_check_call(
4829
+ *args: str,
4830
+ stdout: ta.Any = sys.stderr,
4831
+ input: ta.Any = None, # noqa
4832
+ timeout: ta.Optional[float] = None,
4833
+ **kwargs: ta.Any,
4834
+ ) -> None:
4835
+ _, _ = await _asyncio_subprocess_check_run(
4836
+ *args,
4837
+ stdout=stdout,
4838
+ input=input,
4839
+ timeout=timeout,
4840
+ **kwargs,
4841
+ )
4842
+
4843
+
4844
+ async def asyncio_subprocess_check_output(
4845
+ *args: str,
4846
+ input: ta.Any = None, # noqa
4847
+ timeout: ta.Optional[float] = None,
4848
+ **kwargs: ta.Any,
4849
+ ) -> bytes:
4850
+ stdout, stderr = await _asyncio_subprocess_check_run(
4851
+ *args,
4852
+ stdout=asyncio.subprocess.PIPE,
4853
+ input=input,
4854
+ timeout=timeout,
4855
+ **kwargs,
4856
+ )
4857
+
4858
+ return check_not_none(stdout)
4859
+
4860
+
4861
+ async def asyncio_subprocess_check_output_str(*args: str, **kwargs: ta.Any) -> str:
4862
+ return (await asyncio_subprocess_check_output(*args, **kwargs)).decode().strip()
4863
+
4864
+
4865
+ ##
4866
+
4867
+
4868
+ async def _asyncio_subprocess_try_run(
4869
+ fn: ta.Callable[..., ta.Awaitable[T]],
4870
+ *args: ta.Any,
4871
+ try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
4872
+ **kwargs: ta.Any,
4873
+ ) -> ta.Union[T, Exception]:
4874
+ try:
4875
+ return await fn(*args, **kwargs)
4876
+ except try_exceptions as e: # noqa
4877
+ if log.isEnabledFor(logging.DEBUG):
4878
+ log.exception('command failed')
4879
+ return e
4880
+
4881
+
4882
+ async def asyncio_subprocess_try_call(
4883
+ *args: str,
4884
+ try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
4885
+ **kwargs: ta.Any,
4886
+ ) -> bool:
4887
+ if isinstance(await _asyncio_subprocess_try_run(
4888
+ asyncio_subprocess_check_call,
4889
+ *args,
4890
+ try_exceptions=try_exceptions,
4891
+ **kwargs,
4892
+ ), Exception):
4893
+ return False
4894
+ else:
4895
+ return True
4896
+
4897
+
4898
+ async def asyncio_subprocess_try_output(
4899
+ *args: str,
4900
+ try_exceptions: ta.Tuple[ta.Type[Exception], ...] = DEFAULT_SUBPROCESS_TRY_EXCEPTIONS,
4901
+ **kwargs: ta.Any,
4902
+ ) -> ta.Optional[bytes]:
4903
+ if isinstance(ret := await _asyncio_subprocess_try_run(
4904
+ asyncio_subprocess_check_output,
4905
+ *args,
4906
+ try_exceptions=try_exceptions,
4907
+ **kwargs,
4908
+ ), Exception):
4909
+ return None
4910
+ else:
4911
+ return ret
4912
+
4913
+
4914
+ async def asyncio_subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
4915
+ out = await asyncio_subprocess_try_output(*args, **kwargs)
4916
+ return out.decode().strip() if out is not None else None
4917
+
4918
+
4919
+ ########################################
4920
+ # ../../../omdev/interp/inspect.py
4921
+
4922
+
4923
+ @dc.dataclass(frozen=True)
4924
+ class InterpInspection:
4925
+ exe: str
4926
+ version: Version
4927
+
4928
+ version_str: str
4929
+ config_vars: ta.Mapping[str, str]
4930
+ prefix: str
4931
+ base_prefix: str
4932
+
4933
+ @property
4934
+ def opts(self) -> InterpOpts:
4935
+ return InterpOpts(
4936
+ threaded=bool(self.config_vars.get('Py_GIL_DISABLED')),
4937
+ debug=bool(self.config_vars.get('Py_DEBUG')),
4938
+ )
4939
+
4940
+ @property
4941
+ def iv(self) -> InterpVersion:
4942
+ return InterpVersion(
4943
+ version=self.version,
4944
+ opts=self.opts,
4945
+ )
4946
+
4947
+ @property
4948
+ def is_venv(self) -> bool:
4949
+ return self.prefix != self.base_prefix
4950
+
4951
+
4952
+ class InterpInspector:
4953
+ def __init__(self) -> None:
4954
+ super().__init__()
4955
+
4956
+ self._cache: ta.Dict[str, ta.Optional[InterpInspection]] = {}
4957
+
4958
+ _RAW_INSPECTION_CODE = """
4959
+ __import__('json').dumps(dict(
4960
+ version_str=__import__('sys').version,
4961
+ prefix=__import__('sys').prefix,
4962
+ base_prefix=__import__('sys').base_prefix,
4963
+ config_vars=__import__('sysconfig').get_config_vars(),
4964
+ ))"""
4965
+
4966
+ _INSPECTION_CODE = ''.join(l.strip() for l in _RAW_INSPECTION_CODE.splitlines())
4967
+
4968
+ @staticmethod
4969
+ def _build_inspection(
4970
+ exe: str,
4971
+ output: str,
4972
+ ) -> InterpInspection:
4973
+ dct = json.loads(output)
4974
+
4975
+ version = Version(dct['version_str'].split()[0])
4976
+
4977
+ return InterpInspection(
4978
+ exe=exe,
4979
+ version=version,
4980
+ **{k: dct[k] for k in (
4981
+ 'version_str',
4982
+ 'prefix',
4983
+ 'base_prefix',
4984
+ 'config_vars',
4985
+ )},
4986
+ )
4987
+
4988
+ @classmethod
4989
+ def running(cls) -> 'InterpInspection':
4990
+ return cls._build_inspection(sys.executable, eval(cls._INSPECTION_CODE)) # noqa
4991
+
4992
+ async def _inspect(self, exe: str) -> InterpInspection:
4993
+ output = await asyncio_subprocess_check_output(exe, '-c', f'print({self._INSPECTION_CODE})', quiet=True)
4994
+ return self._build_inspection(exe, output.decode())
4995
+
4996
+ async def inspect(self, exe: str) -> ta.Optional[InterpInspection]:
4997
+ try:
4998
+ return self._cache[exe]
4999
+ except KeyError:
5000
+ ret: ta.Optional[InterpInspection]
5001
+ try:
5002
+ ret = await self._inspect(exe)
5003
+ except Exception as e: # noqa
5004
+ if log.isEnabledFor(logging.DEBUG):
5005
+ log.exception('Failed to inspect interp: %s', exe)
5006
+ ret = None
5007
+ self._cache[exe] = ret
5008
+ return ret
5009
+
5010
+
5011
+ INTERP_INSPECTOR = InterpInspector()
5012
+
5013
+
5014
+ ########################################
5015
+ # ../commands/subprocess.py
5016
+
5017
+
5018
+ ##
5019
+
5020
+
5021
+ @dc.dataclass(frozen=True)
5022
+ class SubprocessCommand(Command['SubprocessCommand.Output']):
5023
+ cmd: ta.Sequence[str]
5024
+
5025
+ shell: bool = False
5026
+ cwd: ta.Optional[str] = None
5027
+ env: ta.Optional[ta.Mapping[str, str]] = None
5028
+
5029
+ stdout: str = 'pipe' # SubprocessChannelOption
5030
+ stderr: str = 'pipe' # SubprocessChannelOption
5031
+
5032
+ input: ta.Optional[bytes] = None
5033
+ timeout: ta.Optional[float] = None
4208
5034
 
4209
5035
  def __post_init__(self) -> None:
4210
5036
  check_not_isinstance(self.cmd, str)
4211
5037
 
4212
- @dc.dataclass(frozen=True)
4213
- class Output(Command.Output):
4214
- rc: int
4215
- pid: int
5038
+ @dc.dataclass(frozen=True)
5039
+ class Output(Command.Output):
5040
+ rc: int
5041
+ pid: int
5042
+
5043
+ elapsed_s: float
5044
+
5045
+ stdout: ta.Optional[bytes] = None
5046
+ stderr: ta.Optional[bytes] = None
5047
+
5048
+
5049
+ ##
5050
+
5051
+
5052
+ class SubprocessCommandExecutor(CommandExecutor[SubprocessCommand, SubprocessCommand.Output]):
5053
+ async def execute(self, cmd: SubprocessCommand) -> SubprocessCommand.Output:
5054
+ proc: asyncio.subprocess.Process
5055
+ async with asyncio_subprocess_popen(
5056
+ *subprocess_maybe_shell_wrap_exec(*cmd.cmd),
5057
+
5058
+ shell=cmd.shell,
5059
+ cwd=cmd.cwd,
5060
+ env={**os.environ, **(cmd.env or {})},
5061
+
5062
+ stdin=subprocess.PIPE if cmd.input is not None else None,
5063
+ stdout=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, cmd.stdout)],
5064
+ stderr=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, cmd.stderr)],
5065
+
5066
+ timeout=cmd.timeout,
5067
+ ) as proc:
5068
+ start_time = time.time()
5069
+ stdout, stderr = await asyncio_subprocess_communicate(
5070
+ proc,
5071
+ input=cmd.input,
5072
+ timeout=cmd.timeout,
5073
+ )
5074
+ end_time = time.time()
5075
+
5076
+ return SubprocessCommand.Output(
5077
+ rc=check_not_none(proc.returncode),
5078
+ pid=proc.pid,
5079
+
5080
+ elapsed_s=end_time - start_time,
5081
+
5082
+ stdout=stdout, # noqa
5083
+ stderr=stderr, # noqa
5084
+ )
5085
+
5086
+
5087
+ ########################################
5088
+ # ../remote/_main.py
5089
+
5090
+
5091
+ ##
5092
+
5093
+
5094
+ class _RemoteExecutionLogHandler(logging.Handler):
5095
+ def __init__(self, fn: ta.Callable[[str], None]) -> None:
5096
+ super().__init__()
5097
+ self._fn = fn
5098
+
5099
+ def emit(self, record):
5100
+ msg = self.format(record)
5101
+ self._fn(msg)
5102
+
5103
+
5104
+ ##
5105
+
5106
+
5107
+ class _RemoteExecutionMain:
5108
+ def __init__(
5109
+ self,
5110
+ chan: RemoteChannel,
5111
+ ) -> None:
5112
+ super().__init__()
5113
+
5114
+ self._chan = chan
5115
+
5116
+ self.__bootstrap: ta.Optional[MainBootstrap] = None
5117
+ self.__injector: ta.Optional[Injector] = None
5118
+
5119
+ @property
5120
+ def _bootstrap(self) -> MainBootstrap:
5121
+ return check_not_none(self.__bootstrap)
5122
+
5123
+ @property
5124
+ def _injector(self) -> Injector:
5125
+ return check_not_none(self.__injector)
5126
+
5127
+ #
5128
+
5129
+ def _timebomb_main(
5130
+ self,
5131
+ delay_s: float,
5132
+ *,
5133
+ sig: int = signal.SIGINT,
5134
+ code: int = 1,
5135
+ ) -> None:
5136
+ time.sleep(delay_s)
5137
+
5138
+ if (pgid := os.getpgid(0)) == os.getpid():
5139
+ os.killpg(pgid, sig)
5140
+
5141
+ os._exit(code) # noqa
5142
+
5143
+ @cached_nullary
5144
+ def _timebomb_thread(self) -> ta.Optional[threading.Thread]:
5145
+ if (tbd := self._bootstrap.remote_config.timebomb_delay_s) is None:
5146
+ return None
5147
+
5148
+ thr = threading.Thread(
5149
+ target=functools.partial(self._timebomb_main, tbd),
5150
+ name=f'{self.__class__.__name__}.timebomb',
5151
+ daemon=True,
5152
+ )
5153
+
5154
+ thr.start()
4216
5155
 
4217
- elapsed_s: float
5156
+ log.debug('Started timebomb thread: %r', thr)
4218
5157
 
4219
- stdout: ta.Optional[bytes] = None
4220
- stderr: ta.Optional[bytes] = None
5158
+ return thr
4221
5159
 
5160
+ #
4222
5161
 
4223
- ##
5162
+ @cached_nullary
5163
+ def _log_handler(self) -> _RemoteLogHandler:
5164
+ return _RemoteLogHandler(self._chan)
4224
5165
 
5166
+ #
4225
5167
 
4226
- class SubprocessCommandExecutor(CommandExecutor[SubprocessCommand, SubprocessCommand.Output]):
4227
- def execute(self, inp: SubprocessCommand) -> SubprocessCommand.Output:
4228
- with subprocess.Popen(
4229
- subprocess_maybe_shell_wrap_exec(*inp.cmd),
5168
+ async def _setup(self) -> None:
5169
+ check_none(self.__bootstrap)
5170
+ check_none(self.__injector)
4230
5171
 
4231
- shell=inp.shell,
4232
- cwd=inp.cwd,
4233
- env={**os.environ, **(inp.env or {})},
5172
+ # Bootstrap
4234
5173
 
4235
- stdin=subprocess.PIPE if inp.input is not None else None,
4236
- stdout=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, inp.stdout)],
4237
- stderr=SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, inp.stderr)],
4238
- ) as proc:
4239
- start_time = time.time()
4240
- stdout, stderr = proc.communicate(
4241
- input=inp.input,
4242
- timeout=inp.timeout,
4243
- )
4244
- end_time = time.time()
5174
+ self.__bootstrap = check_not_none(await self._chan.recv_obj(MainBootstrap))
4245
5175
 
4246
- return SubprocessCommand.Output(
4247
- rc=proc.returncode,
4248
- pid=proc.pid,
5176
+ if (prd := self._bootstrap.remote_config.pycharm_remote_debug) is not None:
5177
+ pycharm_debug_connect(prd)
4249
5178
 
4250
- elapsed_s=end_time - start_time,
5179
+ self.__injector = main_bootstrap(self._bootstrap)
4251
5180
 
4252
- stdout=stdout, # noqa
4253
- stderr=stderr, # noqa
5181
+ self._chan.set_marshaler(self._injector[ObjMarshalerManager])
5182
+
5183
+ # Post-bootstrap
5184
+
5185
+ if self._bootstrap.remote_config.set_pgid:
5186
+ if os.getpgid(0) != os.getpid():
5187
+ log.debug('Setting pgid')
5188
+ os.setpgid(0, 0)
5189
+
5190
+ if (ds := self._bootstrap.remote_config.deathsig) is not None:
5191
+ log.debug('Setting deathsig: %s', ds)
5192
+ set_process_deathsig(int(signal.Signals[f'SIG{ds.upper()}']))
5193
+
5194
+ self._timebomb_thread()
5195
+
5196
+ if self._bootstrap.remote_config.forward_logging:
5197
+ log.debug('Installing log forwarder')
5198
+ logging.root.addHandler(self._log_handler())
5199
+
5200
+ #
5201
+
5202
+ async def run(self) -> None:
5203
+ await self._setup()
5204
+
5205
+ executor = self._injector[LocalCommandExecutor]
5206
+
5207
+ handler = _RemoteCommandHandler(self._chan, executor)
5208
+
5209
+ await handler.run()
5210
+
5211
+
5212
+ def _remote_execution_main() -> None:
5213
+ rt = pyremote_bootstrap_finalize() # noqa
5214
+
5215
+ async def inner() -> None:
5216
+ input = await asyncio_open_stream_reader(rt.input) # noqa
5217
+ output = await asyncio_open_stream_writer(rt.output)
5218
+
5219
+ chan = RemoteChannelImpl(
5220
+ input,
5221
+ output,
4254
5222
  )
4255
5223
 
5224
+ await _RemoteExecutionMain(chan).run()
5225
+
5226
+ asyncio.run(inner())
5227
+
4256
5228
 
4257
5229
  ########################################
4258
5230
  # ../remote/spawning.py
4259
5231
 
4260
5232
 
4261
- class RemoteSpawning:
5233
+ ##
5234
+
5235
+
5236
+ class RemoteSpawning(abc.ABC):
4262
5237
  @dc.dataclass(frozen=True)
4263
5238
  class Target:
4264
5239
  shell: ta.Optional[str] = None
@@ -4269,15 +5244,35 @@ class RemoteSpawning:
4269
5244
 
4270
5245
  stderr: ta.Optional[str] = None # SubprocessChannelOption
4271
5246
 
4272
- #
5247
+ @dc.dataclass(frozen=True)
5248
+ class Spawned:
5249
+ stdin: asyncio.StreamWriter
5250
+ stdout: asyncio.StreamReader
5251
+ stderr: ta.Optional[asyncio.StreamReader]
5252
+
5253
+ @abc.abstractmethod
5254
+ def spawn(
5255
+ self,
5256
+ tgt: Target,
5257
+ src: str,
5258
+ *,
5259
+ timeout: ta.Optional[float] = None,
5260
+ debug: bool = False,
5261
+ ) -> ta.AsyncContextManager[Spawned]:
5262
+ raise NotImplementedError
4273
5263
 
4274
- class _PreparedCmd(ta.NamedTuple):
5264
+
5265
+ ##
5266
+
5267
+
5268
+ class SubprocessRemoteSpawning(RemoteSpawning):
5269
+ class _PreparedCmd(ta.NamedTuple): # noqa
4275
5270
  cmd: ta.Sequence[str]
4276
5271
  shell: bool
4277
5272
 
4278
5273
  def _prepare_cmd(
4279
5274
  self,
4280
- tgt: Target,
5275
+ tgt: RemoteSpawning.Target,
4281
5276
  src: str,
4282
5277
  ) -> _PreparedCmd:
4283
5278
  if tgt.shell is not None:
@@ -4285,44 +5280,38 @@ class RemoteSpawning:
4285
5280
  if tgt.shell_quote:
4286
5281
  sh_src = shlex.quote(sh_src)
4287
5282
  sh_cmd = f'{tgt.shell} {sh_src}'
4288
- return RemoteSpawning._PreparedCmd(
4289
- cmd=[sh_cmd],
4290
- shell=True,
4291
- )
5283
+ return SubprocessRemoteSpawning._PreparedCmd([sh_cmd], shell=True)
4292
5284
 
4293
5285
  else:
4294
- return RemoteSpawning._PreparedCmd(
4295
- cmd=[tgt.python, '-c', src],
4296
- shell=False,
4297
- )
5286
+ return SubprocessRemoteSpawning._PreparedCmd([tgt.python, '-c', src], shell=False)
4298
5287
 
4299
5288
  #
4300
5289
 
4301
- @dc.dataclass(frozen=True)
4302
- class Spawned:
4303
- stdin: ta.IO
4304
- stdout: ta.IO
4305
- stderr: ta.Optional[ta.IO]
4306
-
4307
- @contextlib.contextmanager
4308
- def spawn(
5290
+ @contextlib.asynccontextmanager
5291
+ async def spawn(
4309
5292
  self,
4310
- tgt: Target,
5293
+ tgt: RemoteSpawning.Target,
4311
5294
  src: str,
4312
5295
  *,
4313
5296
  timeout: ta.Optional[float] = None,
4314
- ) -> ta.Generator[Spawned, None, None]:
5297
+ debug: bool = False,
5298
+ ) -> ta.AsyncGenerator[RemoteSpawning.Spawned, None]:
4315
5299
  pc = self._prepare_cmd(tgt, src)
4316
5300
 
4317
- with subprocess.Popen(
4318
- subprocess_maybe_shell_wrap_exec(*pc.cmd),
4319
- shell=pc.shell,
4320
- stdin=subprocess.PIPE,
4321
- stdout=subprocess.PIPE,
4322
- stderr=(
4323
- SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, tgt.stderr)]
4324
- if tgt.stderr is not None else None
4325
- ),
5301
+ cmd = pc.cmd
5302
+ if not debug:
5303
+ cmd = subprocess_maybe_shell_wrap_exec(*cmd)
5304
+
5305
+ async with asyncio_subprocess_popen(
5306
+ *cmd,
5307
+ shell=pc.shell,
5308
+ stdin=subprocess.PIPE,
5309
+ stdout=subprocess.PIPE,
5310
+ stderr=(
5311
+ SUBPROCESS_CHANNEL_OPTION_VALUES[ta.cast(SubprocessChannelOption, tgt.stderr)]
5312
+ if tgt.stderr is not None else None
5313
+ ),
5314
+ timeout=timeout,
4326
5315
  ) as proc:
4327
5316
  stdin = check_not_none(proc.stdin)
4328
5317
  stdout = check_not_none(proc.stdout)
@@ -4340,8 +5329,6 @@ class RemoteSpawning:
4340
5329
  except BrokenPipeError:
4341
5330
  pass
4342
5331
 
4343
- proc.wait(timeout)
4344
-
4345
5332
 
4346
5333
  ########################################
4347
5334
  # ../../../omdev/interp/providers.py
@@ -4370,17 +5357,17 @@ class InterpProvider(abc.ABC):
4370
5357
  setattr(cls, 'name', snake_case(cls.__name__[:-len(sfx)]))
4371
5358
 
4372
5359
  @abc.abstractmethod
4373
- def get_installed_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5360
+ def get_installed_versions(self, spec: InterpSpecifier) -> ta.Awaitable[ta.Sequence[InterpVersion]]:
4374
5361
  raise NotImplementedError
4375
5362
 
4376
5363
  @abc.abstractmethod
4377
- def get_installed_version(self, version: InterpVersion) -> Interp:
5364
+ def get_installed_version(self, version: InterpVersion) -> ta.Awaitable[Interp]:
4378
5365
  raise NotImplementedError
4379
5366
 
4380
- def get_installable_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5367
+ async def get_installable_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
4381
5368
  return []
4382
5369
 
4383
- def install_version(self, version: InterpVersion) -> Interp:
5370
+ async def install_version(self, version: InterpVersion) -> Interp:
4384
5371
  raise TypeError
4385
5372
 
4386
5373
 
@@ -4392,10 +5379,10 @@ class RunningInterpProvider(InterpProvider):
4392
5379
  def version(self) -> InterpVersion:
4393
5380
  return InterpInspector.running().iv
4394
5381
 
4395
- def get_installed_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5382
+ async def get_installed_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
4396
5383
  return [self.version()]
4397
5384
 
4398
- def get_installed_version(self, version: InterpVersion) -> Interp:
5385
+ async def get_installed_version(self, version: InterpVersion) -> Interp:
4399
5386
  if version != self.version():
4400
5387
  raise KeyError(version)
4401
5388
  return Interp(
@@ -4405,159 +5392,26 @@ class RunningInterpProvider(InterpProvider):
4405
5392
 
4406
5393
 
4407
5394
  ########################################
4408
- # ../remote/execution.py
4409
-
4410
-
4411
- ##
4412
-
4413
-
4414
- class _RemoteExecutionLogHandler(logging.Handler):
4415
- def __init__(self, fn: ta.Callable[[str], None]) -> None:
4416
- super().__init__()
4417
- self._fn = fn
4418
-
4419
- def emit(self, record):
4420
- msg = self.format(record)
4421
- self._fn(msg)
4422
-
4423
-
4424
- @dc.dataclass(frozen=True)
4425
- class _RemoteExecutionRequest:
4426
- c: Command
4427
-
4428
-
4429
- @dc.dataclass(frozen=True)
4430
- class _RemoteExecutionLog:
4431
- s: str
4432
-
4433
-
4434
- @dc.dataclass(frozen=True)
4435
- class _RemoteExecutionResponse:
4436
- r: ta.Optional[CommandOutputOrExceptionData] = None
4437
- l: ta.Optional[_RemoteExecutionLog] = None
4438
-
4439
-
4440
- def _remote_execution_main() -> None:
4441
- rt = pyremote_bootstrap_finalize() # noqa
4442
-
4443
- chan = RemoteChannel(
4444
- rt.input,
4445
- rt.output,
4446
- )
4447
-
4448
- bs = check_not_none(chan.recv_obj(MainBootstrap))
4449
-
4450
- if (prd := bs.remote_config.pycharm_remote_debug) is not None:
4451
- pycharm_debug_connect(prd)
4452
-
4453
- injector = main_bootstrap(bs)
4454
-
4455
- chan.set_marshaler(injector[ObjMarshalerManager])
4456
-
4457
- #
4458
-
4459
- log_lock = threading.RLock()
4460
- send_logs = False
4461
-
4462
- def log_fn(s: str) -> None:
4463
- with log_lock:
4464
- if send_logs:
4465
- chan.send_obj(_RemoteExecutionResponse(l=_RemoteExecutionLog(s)))
4466
-
4467
- log_handler = _RemoteExecutionLogHandler(log_fn)
4468
- logging.root.addHandler(log_handler)
4469
-
4470
- #
4471
-
4472
- ce = injector[LocalCommandExecutor]
4473
-
4474
- while True:
4475
- req = chan.recv_obj(_RemoteExecutionRequest)
4476
- if req is None:
4477
- break
4478
-
4479
- with log_lock:
4480
- send_logs = True
4481
-
4482
- r = ce.try_execute(
4483
- req.c,
4484
- log=log,
4485
- omit_exc_object=True,
4486
- )
4487
-
4488
- with log_lock:
4489
- send_logs = False
4490
-
4491
- chan.send_obj(_RemoteExecutionResponse(r=CommandOutputOrExceptionData(
4492
- output=r.output,
4493
- exception=r.exception,
4494
- )))
5395
+ # ../remote/connection.py
4495
5396
 
4496
5397
 
4497
5398
  ##
4498
5399
 
4499
5400
 
4500
- @dc.dataclass()
4501
- class RemoteCommandError(Exception):
4502
- e: CommandException
4503
-
4504
-
4505
- class RemoteCommandExecutor(CommandExecutor):
4506
- def __init__(self, chan: RemoteChannel) -> None:
4507
- super().__init__()
4508
-
4509
- self._chan = chan
4510
-
4511
- def _remote_execute(self, cmd: Command) -> CommandOutputOrException:
4512
- self._chan.send_obj(_RemoteExecutionRequest(cmd))
4513
-
4514
- while True:
4515
- if (r := self._chan.recv_obj(_RemoteExecutionResponse)) is None:
4516
- raise EOFError
4517
-
4518
- if r.l is not None:
4519
- log.info(r.l.s)
4520
-
4521
- if r.r is not None:
4522
- return r.r
4523
-
4524
- # @ta.override
4525
- def execute(self, cmd: Command) -> Command.Output:
4526
- r = self._remote_execute(cmd)
4527
- if (e := r.exception) is not None:
4528
- raise RemoteCommandError(e)
4529
- else:
4530
- return check_not_none(r.output)
4531
-
4532
- # @ta.override
4533
- def try_execute(
5401
+ class RemoteExecutionConnector(abc.ABC):
5402
+ @abc.abstractmethod
5403
+ def connect(
4534
5404
  self,
4535
- cmd: Command,
4536
- *,
4537
- log: ta.Optional[logging.Logger] = None,
4538
- omit_exc_object: bool = False,
4539
- ) -> CommandOutputOrException:
4540
- try:
4541
- r = self._remote_execute(cmd)
4542
-
4543
- except Exception as e: # noqa
4544
- if log is not None:
4545
- log.exception('Exception executing remote command: %r', type(cmd))
4546
-
4547
- return CommandOutputOrExceptionData(exception=CommandException.of(
4548
- e,
4549
- omit_exc_object=omit_exc_object,
4550
- cmd=cmd,
4551
- ))
4552
-
4553
- else:
4554
- return r
5405
+ tgt: RemoteSpawning.Target,
5406
+ bs: MainBootstrap,
5407
+ ) -> ta.AsyncContextManager[RemoteCommandExecutor]:
5408
+ raise NotImplementedError
4555
5409
 
4556
5410
 
4557
5411
  ##
4558
5412
 
4559
5413
 
4560
- class RemoteExecution:
5414
+ class PyremoteRemoteExecutionConnector(RemoteExecutionConnector):
4561
5415
  def __init__(
4562
5416
  self,
4563
5417
  *,
@@ -4590,38 +5444,43 @@ class RemoteExecution:
4590
5444
 
4591
5445
  #
4592
5446
 
4593
- @contextlib.contextmanager
4594
- def connect(
5447
+ @contextlib.asynccontextmanager
5448
+ async def connect(
4595
5449
  self,
4596
5450
  tgt: RemoteSpawning.Target,
4597
5451
  bs: MainBootstrap,
4598
- ) -> ta.Generator[RemoteCommandExecutor, None, None]:
5452
+ ) -> ta.AsyncGenerator[RemoteCommandExecutor, None]:
4599
5453
  spawn_src = self._spawn_src()
4600
5454
  remote_src = self._remote_src()
4601
5455
 
4602
- with self._spawning.spawn(
5456
+ async with self._spawning.spawn(
4603
5457
  tgt,
4604
5458
  spawn_src,
5459
+ debug=bs.main_config.debug,
4605
5460
  ) as proc:
4606
- res = PyremoteBootstrapDriver( # noqa
5461
+ res = await PyremoteBootstrapDriver( # noqa
4607
5462
  remote_src,
4608
5463
  PyremoteBootstrapOptions(
4609
5464
  debug=bs.main_config.debug,
4610
5465
  ),
4611
- ).run(
5466
+ ).async_run(
4612
5467
  proc.stdout,
4613
5468
  proc.stdin,
4614
5469
  )
4615
5470
 
4616
- chan = RemoteChannel(
5471
+ chan = RemoteChannelImpl(
4617
5472
  proc.stdout,
4618
5473
  proc.stdin,
4619
5474
  msh=self._msh,
4620
5475
  )
4621
5476
 
4622
- chan.send_obj(bs)
5477
+ await chan.send_obj(bs)
5478
+
5479
+ rce: RemoteCommandExecutor
5480
+ async with contextlib.aclosing(RemoteCommandExecutor(chan)) as rce:
5481
+ await rce.start()
4623
5482
 
4624
- yield RemoteCommandExecutor(chan)
5483
+ yield rce
4625
5484
 
4626
5485
 
4627
5486
  ########################################
@@ -4643,7 +5502,6 @@ TODO:
4643
5502
 
4644
5503
 
4645
5504
  class Pyenv:
4646
-
4647
5505
  def __init__(
4648
5506
  self,
4649
5507
  *,
@@ -4656,13 +5514,13 @@ class Pyenv:
4656
5514
 
4657
5515
  self._root_kw = root
4658
5516
 
4659
- @cached_nullary
4660
- def root(self) -> ta.Optional[str]:
5517
+ @async_cached_nullary
5518
+ async def root(self) -> ta.Optional[str]:
4661
5519
  if self._root_kw is not None:
4662
5520
  return self._root_kw
4663
5521
 
4664
5522
  if shutil.which('pyenv'):
4665
- return subprocess_check_output_str('pyenv', 'root')
5523
+ return await asyncio_subprocess_check_output_str('pyenv', 'root')
4666
5524
 
4667
5525
  d = os.path.expanduser('~/.pyenv')
4668
5526
  if os.path.isdir(d) and os.path.isfile(os.path.join(d, 'bin', 'pyenv')):
@@ -4670,12 +5528,12 @@ class Pyenv:
4670
5528
 
4671
5529
  return None
4672
5530
 
4673
- @cached_nullary
4674
- def exe(self) -> str:
4675
- return os.path.join(check_not_none(self.root()), 'bin', 'pyenv')
5531
+ @async_cached_nullary
5532
+ async def exe(self) -> str:
5533
+ return os.path.join(check_not_none(await self.root()), 'bin', 'pyenv')
4676
5534
 
4677
- def version_exes(self) -> ta.List[ta.Tuple[str, str]]:
4678
- if (root := self.root()) is None:
5535
+ async def version_exes(self) -> ta.List[ta.Tuple[str, str]]:
5536
+ if (root := await self.root()) is None:
4679
5537
  return []
4680
5538
  ret = []
4681
5539
  vp = os.path.join(root, 'versions')
@@ -4687,11 +5545,11 @@ class Pyenv:
4687
5545
  ret.append((dn, ep))
4688
5546
  return ret
4689
5547
 
4690
- def installable_versions(self) -> ta.List[str]:
4691
- if self.root() is None:
5548
+ async def installable_versions(self) -> ta.List[str]:
5549
+ if await self.root() is None:
4692
5550
  return []
4693
5551
  ret = []
4694
- s = subprocess_check_output_str(self.exe(), 'install', '--list')
5552
+ s = await asyncio_subprocess_check_output_str(await self.exe(), 'install', '--list')
4695
5553
  for l in s.splitlines():
4696
5554
  if not l.startswith(' '):
4697
5555
  continue
@@ -4701,12 +5559,12 @@ class Pyenv:
4701
5559
  ret.append(l)
4702
5560
  return ret
4703
5561
 
4704
- def update(self) -> bool:
4705
- if (root := self.root()) is None:
5562
+ async def update(self) -> bool:
5563
+ if (root := await self.root()) is None:
4706
5564
  return False
4707
5565
  if not os.path.isdir(os.path.join(root, '.git')):
4708
5566
  return False
4709
- subprocess_check_call('git', 'pull', cwd=root)
5567
+ await asyncio_subprocess_check_call('git', 'pull', cwd=root)
4710
5568
  return True
4711
5569
 
4712
5570
 
@@ -4767,17 +5625,16 @@ THREADED_PYENV_INSTALL_OPTS = PyenvInstallOpts(conf_opts=['--disable-gil'])
4767
5625
 
4768
5626
  class PyenvInstallOptsProvider(abc.ABC):
4769
5627
  @abc.abstractmethod
4770
- def opts(self) -> PyenvInstallOpts:
5628
+ def opts(self) -> ta.Awaitable[PyenvInstallOpts]:
4771
5629
  raise NotImplementedError
4772
5630
 
4773
5631
 
4774
5632
  class LinuxPyenvInstallOpts(PyenvInstallOptsProvider):
4775
- def opts(self) -> PyenvInstallOpts:
5633
+ async def opts(self) -> PyenvInstallOpts:
4776
5634
  return PyenvInstallOpts()
4777
5635
 
4778
5636
 
4779
5637
  class DarwinPyenvInstallOpts(PyenvInstallOptsProvider):
4780
-
4781
5638
  @cached_nullary
4782
5639
  def framework_opts(self) -> PyenvInstallOpts:
4783
5640
  return PyenvInstallOpts(conf_opts=['--enable-framework'])
@@ -4793,12 +5650,12 @@ class DarwinPyenvInstallOpts(PyenvInstallOptsProvider):
4793
5650
  'zlib',
4794
5651
  ]
4795
5652
 
4796
- @cached_nullary
4797
- def brew_deps_opts(self) -> PyenvInstallOpts:
5653
+ @async_cached_nullary
5654
+ async def brew_deps_opts(self) -> PyenvInstallOpts:
4798
5655
  cflags = []
4799
5656
  ldflags = []
4800
5657
  for dep in self.BREW_DEPS:
4801
- dep_prefix = subprocess_check_output_str('brew', '--prefix', dep)
5658
+ dep_prefix = await asyncio_subprocess_check_output_str('brew', '--prefix', dep)
4802
5659
  cflags.append(f'-I{dep_prefix}/include')
4803
5660
  ldflags.append(f'-L{dep_prefix}/lib')
4804
5661
  return PyenvInstallOpts(
@@ -4806,13 +5663,13 @@ class DarwinPyenvInstallOpts(PyenvInstallOptsProvider):
4806
5663
  ldflags=ldflags,
4807
5664
  )
4808
5665
 
4809
- @cached_nullary
4810
- def brew_tcl_opts(self) -> PyenvInstallOpts:
4811
- if subprocess_try_output('brew', '--prefix', 'tcl-tk') is None:
5666
+ @async_cached_nullary
5667
+ async def brew_tcl_opts(self) -> PyenvInstallOpts:
5668
+ if await asyncio_subprocess_try_output('brew', '--prefix', 'tcl-tk') is None:
4812
5669
  return PyenvInstallOpts()
4813
5670
 
4814
- tcl_tk_prefix = subprocess_check_output_str('brew', '--prefix', 'tcl-tk')
4815
- tcl_tk_ver_str = subprocess_check_output_str('brew', 'ls', '--versions', 'tcl-tk')
5671
+ tcl_tk_prefix = await asyncio_subprocess_check_output_str('brew', '--prefix', 'tcl-tk')
5672
+ tcl_tk_ver_str = await asyncio_subprocess_check_output_str('brew', 'ls', '--versions', 'tcl-tk')
4816
5673
  tcl_tk_ver = '.'.join(tcl_tk_ver_str.split()[1].split('.')[:2])
4817
5674
 
4818
5675
  return PyenvInstallOpts(conf_opts=[
@@ -4827,11 +5684,11 @@ class DarwinPyenvInstallOpts(PyenvInstallOptsProvider):
4827
5684
  # pkg_config_path += ':' + os.environ['PKG_CONFIG_PATH']
4828
5685
  # return PyenvInstallOpts(env={'PKG_CONFIG_PATH': pkg_config_path})
4829
5686
 
4830
- def opts(self) -> PyenvInstallOpts:
5687
+ async def opts(self) -> PyenvInstallOpts:
4831
5688
  return PyenvInstallOpts().merge(
4832
5689
  self.framework_opts(),
4833
- self.brew_deps_opts(),
4834
- self.brew_tcl_opts(),
5690
+ await self.brew_deps_opts(),
5691
+ await self.brew_tcl_opts(),
4835
5692
  # self.brew_ssl_opts(),
4836
5693
  )
4837
5694
 
@@ -4863,20 +5720,8 @@ class PyenvVersionInstaller:
4863
5720
  ) -> None:
4864
5721
  super().__init__()
4865
5722
 
4866
- if no_default_opts:
4867
- if opts is None:
4868
- opts = PyenvInstallOpts()
4869
- else:
4870
- lst = [opts if opts is not None else DEFAULT_PYENV_INSTALL_OPTS]
4871
- if interp_opts.debug:
4872
- lst.append(DEBUG_PYENV_INSTALL_OPTS)
4873
- if interp_opts.threaded:
4874
- lst.append(THREADED_PYENV_INSTALL_OPTS)
4875
- lst.append(PLATFORM_PYENV_INSTALL_OPTS[sys.platform].opts())
4876
- opts = PyenvInstallOpts().merge(*lst)
4877
-
4878
5723
  self._version = version
4879
- self._opts = opts
5724
+ self._given_opts = opts
4880
5725
  self._interp_opts = interp_opts
4881
5726
  self._given_install_name = install_name
4882
5727
 
@@ -4887,9 +5732,21 @@ class PyenvVersionInstaller:
4887
5732
  def version(self) -> str:
4888
5733
  return self._version
4889
5734
 
4890
- @property
4891
- def opts(self) -> PyenvInstallOpts:
4892
- return self._opts
5735
+ @async_cached_nullary
5736
+ async def opts(self) -> PyenvInstallOpts:
5737
+ opts = self._given_opts
5738
+ if self._no_default_opts:
5739
+ if opts is None:
5740
+ opts = PyenvInstallOpts()
5741
+ else:
5742
+ lst = [self._given_opts if self._given_opts is not None else DEFAULT_PYENV_INSTALL_OPTS]
5743
+ if self._interp_opts.debug:
5744
+ lst.append(DEBUG_PYENV_INSTALL_OPTS)
5745
+ if self._interp_opts.threaded:
5746
+ lst.append(THREADED_PYENV_INSTALL_OPTS)
5747
+ lst.append(await PLATFORM_PYENV_INSTALL_OPTS[sys.platform].opts())
5748
+ opts = PyenvInstallOpts().merge(*lst)
5749
+ return opts
4893
5750
 
4894
5751
  @cached_nullary
4895
5752
  def install_name(self) -> str:
@@ -4897,17 +5754,18 @@ class PyenvVersionInstaller:
4897
5754
  return self._given_install_name
4898
5755
  return self._version + ('-debug' if self._interp_opts.debug else '')
4899
5756
 
4900
- @cached_nullary
4901
- def install_dir(self) -> str:
4902
- return str(os.path.join(check_not_none(self._pyenv.root()), 'versions', self.install_name()))
5757
+ @async_cached_nullary
5758
+ async def install_dir(self) -> str:
5759
+ return str(os.path.join(check_not_none(await self._pyenv.root()), 'versions', self.install_name()))
4903
5760
 
4904
- @cached_nullary
4905
- def install(self) -> str:
4906
- env = {**os.environ, **self._opts.env}
5761
+ @async_cached_nullary
5762
+ async def install(self) -> str:
5763
+ opts = await self.opts()
5764
+ env = {**os.environ, **opts.env}
4907
5765
  for k, l in [
4908
- ('CFLAGS', self._opts.cflags),
4909
- ('LDFLAGS', self._opts.ldflags),
4910
- ('PYTHON_CONFIGURE_OPTS', self._opts.conf_opts),
5766
+ ('CFLAGS', opts.cflags),
5767
+ ('LDFLAGS', opts.ldflags),
5768
+ ('PYTHON_CONFIGURE_OPTS', opts.conf_opts),
4911
5769
  ]:
4912
5770
  v = ' '.join(l)
4913
5771
  if k in os.environ:
@@ -4915,13 +5773,13 @@ class PyenvVersionInstaller:
4915
5773
  env[k] = v
4916
5774
 
4917
5775
  conf_args = [
4918
- *self._opts.opts,
5776
+ *opts.opts,
4919
5777
  self._version,
4920
5778
  ]
4921
5779
 
4922
5780
  if self._given_install_name is not None:
4923
5781
  full_args = [
4924
- os.path.join(check_not_none(self._pyenv.root()), 'plugins', 'python-build', 'bin', 'python-build'),
5782
+ os.path.join(check_not_none(await self._pyenv.root()), 'plugins', 'python-build', 'bin', 'python-build'), # noqa
4925
5783
  *conf_args,
4926
5784
  self.install_dir(),
4927
5785
  ]
@@ -4932,12 +5790,12 @@ class PyenvVersionInstaller:
4932
5790
  *conf_args,
4933
5791
  ]
4934
5792
 
4935
- subprocess_check_call(
5793
+ await asyncio_subprocess_check_call(
4936
5794
  *full_args,
4937
5795
  env=env,
4938
5796
  )
4939
5797
 
4940
- exe = os.path.join(self.install_dir(), 'bin', 'python')
5798
+ exe = os.path.join(await self.install_dir(), 'bin', 'python')
4941
5799
  if not os.path.isfile(exe):
4942
5800
  raise RuntimeError(f'Interpreter not found: {exe}')
4943
5801
  return exe
@@ -4947,7 +5805,6 @@ class PyenvVersionInstaller:
4947
5805
 
4948
5806
 
4949
5807
  class PyenvInterpProvider(InterpProvider):
4950
-
4951
5808
  def __init__(
4952
5809
  self,
4953
5810
  pyenv: Pyenv = Pyenv(),
@@ -4990,11 +5847,11 @@ class PyenvInterpProvider(InterpProvider):
4990
5847
  exe: str
4991
5848
  version: InterpVersion
4992
5849
 
4993
- def _make_installed(self, vn: str, ep: str) -> ta.Optional[Installed]:
5850
+ async def _make_installed(self, vn: str, ep: str) -> ta.Optional[Installed]:
4994
5851
  iv: ta.Optional[InterpVersion]
4995
5852
  if self._inspect:
4996
5853
  try:
4997
- iv = check_not_none(self._inspector.inspect(ep)).iv
5854
+ iv = check_not_none(await self._inspector.inspect(ep)).iv
4998
5855
  except Exception as e: # noqa
4999
5856
  return None
5000
5857
  else:
@@ -5007,10 +5864,10 @@ class PyenvInterpProvider(InterpProvider):
5007
5864
  version=iv,
5008
5865
  )
5009
5866
 
5010
- def installed(self) -> ta.Sequence[Installed]:
5867
+ async def installed(self) -> ta.Sequence[Installed]:
5011
5868
  ret: ta.List[PyenvInterpProvider.Installed] = []
5012
- for vn, ep in self._pyenv.version_exes():
5013
- if (i := self._make_installed(vn, ep)) is None:
5869
+ for vn, ep in await self._pyenv.version_exes():
5870
+ if (i := await self._make_installed(vn, ep)) is None:
5014
5871
  log.debug('Invalid pyenv version: %s', vn)
5015
5872
  continue
5016
5873
  ret.append(i)
@@ -5018,11 +5875,11 @@ class PyenvInterpProvider(InterpProvider):
5018
5875
 
5019
5876
  #
5020
5877
 
5021
- def get_installed_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5022
- return [i.version for i in self.installed()]
5878
+ async def get_installed_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5879
+ return [i.version for i in await self.installed()]
5023
5880
 
5024
- def get_installed_version(self, version: InterpVersion) -> Interp:
5025
- for i in self.installed():
5881
+ async def get_installed_version(self, version: InterpVersion) -> Interp:
5882
+ for i in await self.installed():
5026
5883
  if i.version == version:
5027
5884
  return Interp(
5028
5885
  exe=i.exe,
@@ -5032,10 +5889,10 @@ class PyenvInterpProvider(InterpProvider):
5032
5889
 
5033
5890
  #
5034
5891
 
5035
- def _get_installable_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5892
+ async def _get_installable_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5036
5893
  lst = []
5037
5894
 
5038
- for vs in self._pyenv.installable_versions():
5895
+ for vs in await self._pyenv.installable_versions():
5039
5896
  if (iv := self.guess_version(vs)) is None:
5040
5897
  continue
5041
5898
  if iv.opts.debug:
@@ -5045,16 +5902,16 @@ class PyenvInterpProvider(InterpProvider):
5045
5902
 
5046
5903
  return lst
5047
5904
 
5048
- def get_installable_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5049
- lst = self._get_installable_versions(spec)
5905
+ async def get_installable_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5906
+ lst = await self._get_installable_versions(spec)
5050
5907
 
5051
5908
  if self._try_update and not any(v in spec for v in lst):
5052
5909
  if self._pyenv.update():
5053
- lst = self._get_installable_versions(spec)
5910
+ lst = await self._get_installable_versions(spec)
5054
5911
 
5055
5912
  return lst
5056
5913
 
5057
- def install_version(self, version: InterpVersion) -> Interp:
5914
+ async def install_version(self, version: InterpVersion) -> Interp:
5058
5915
  inst_version = str(version.version)
5059
5916
  inst_opts = version.opts
5060
5917
  if inst_opts.threaded:
@@ -5066,7 +5923,7 @@ class PyenvInterpProvider(InterpProvider):
5066
5923
  interp_opts=inst_opts,
5067
5924
  )
5068
5925
 
5069
- exe = installer.install()
5926
+ exe = await installer.install()
5070
5927
  return Interp(exe, version)
5071
5928
 
5072
5929
 
@@ -5145,7 +6002,7 @@ class SystemInterpProvider(InterpProvider):
5145
6002
 
5146
6003
  #
5147
6004
 
5148
- def get_exe_version(self, exe: str) -> ta.Optional[InterpVersion]:
6005
+ async def get_exe_version(self, exe: str) -> ta.Optional[InterpVersion]:
5149
6006
  if not self.inspect:
5150
6007
  s = os.path.basename(exe)
5151
6008
  if s.startswith('python'):
@@ -5155,13 +6012,13 @@ class SystemInterpProvider(InterpProvider):
5155
6012
  return InterpVersion.parse(s)
5156
6013
  except InvalidVersion:
5157
6014
  pass
5158
- ii = self.inspector.inspect(exe)
6015
+ ii = await self.inspector.inspect(exe)
5159
6016
  return ii.iv if ii is not None else None
5160
6017
 
5161
- def exe_versions(self) -> ta.Sequence[ta.Tuple[str, InterpVersion]]:
6018
+ async def exe_versions(self) -> ta.Sequence[ta.Tuple[str, InterpVersion]]:
5162
6019
  lst = []
5163
6020
  for e in self.exes():
5164
- if (ev := self.get_exe_version(e)) is None:
6021
+ if (ev := await self.get_exe_version(e)) is None:
5165
6022
  log.debug('Invalid system version: %s', e)
5166
6023
  continue
5167
6024
  lst.append((e, ev))
@@ -5169,11 +6026,11 @@ class SystemInterpProvider(InterpProvider):
5169
6026
 
5170
6027
  #
5171
6028
 
5172
- def get_installed_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
5173
- return [ev for e, ev in self.exe_versions()]
6029
+ async def get_installed_versions(self, spec: InterpSpecifier) -> ta.Sequence[InterpVersion]:
6030
+ return [ev for e, ev in await self.exe_versions()]
5174
6031
 
5175
- def get_installed_version(self, version: InterpVersion) -> Interp:
5176
- for e, ev in self.exe_versions():
6032
+ async def get_installed_version(self, version: InterpVersion) -> Interp:
6033
+ for e, ev in await self.exe_versions():
5177
6034
  if ev != version:
5178
6035
  continue
5179
6036
  return Interp(
@@ -5194,9 +6051,11 @@ def bind_remote(
5194
6051
  lst: ta.List[InjectorBindingOrBindings] = [
5195
6052
  inj.bind(remote_config),
5196
6053
 
5197
- inj.bind(RemoteSpawning, singleton=True),
6054
+ inj.bind(SubprocessRemoteSpawning, singleton=True),
6055
+ inj.bind(RemoteSpawning, to_key=SubprocessRemoteSpawning),
5198
6056
 
5199
- inj.bind(RemoteExecution, singleton=True),
6057
+ inj.bind(PyremoteRemoteExecutionConnector, singleton=True),
6058
+ inj.bind(RemoteExecutionConnector, to_key=PyremoteRemoteExecutionConnector),
5200
6059
  ]
5201
6060
 
5202
6061
  if (pf := remote_config.payload_file) is not None:
@@ -5220,13 +6079,14 @@ class InterpResolver:
5220
6079
  providers: ta.Sequence[ta.Tuple[str, InterpProvider]],
5221
6080
  ) -> None:
5222
6081
  super().__init__()
6082
+
5223
6083
  self._providers: ta.Mapping[str, InterpProvider] = collections.OrderedDict(providers)
5224
6084
 
5225
- def _resolve_installed(self, spec: InterpSpecifier) -> ta.Optional[ta.Tuple[InterpProvider, InterpVersion]]:
6085
+ async def _resolve_installed(self, spec: InterpSpecifier) -> ta.Optional[ta.Tuple[InterpProvider, InterpVersion]]:
5226
6086
  lst = [
5227
6087
  (i, si)
5228
6088
  for i, p in enumerate(self._providers.values())
5229
- for si in p.get_installed_versions(spec)
6089
+ for si in await p.get_installed_versions(spec)
5230
6090
  if spec.contains(si)
5231
6091
  ]
5232
6092
 
@@ -5238,16 +6098,16 @@ class InterpResolver:
5238
6098
  bp = list(self._providers.values())[bi]
5239
6099
  return (bp, bv)
5240
6100
 
5241
- def resolve(
6101
+ async def resolve(
5242
6102
  self,
5243
6103
  spec: InterpSpecifier,
5244
6104
  *,
5245
6105
  install: bool = False,
5246
6106
  ) -> ta.Optional[Interp]:
5247
- tup = self._resolve_installed(spec)
6107
+ tup = await self._resolve_installed(spec)
5248
6108
  if tup is not None:
5249
6109
  bp, bv = tup
5250
- return bp.get_installed_version(bv)
6110
+ return await bp.get_installed_version(bv)
5251
6111
 
5252
6112
  if not install:
5253
6113
  return None
@@ -5255,21 +6115,21 @@ class InterpResolver:
5255
6115
  tp = list(self._providers.values())[0] # noqa
5256
6116
 
5257
6117
  sv = sorted(
5258
- [s for s in tp.get_installable_versions(spec) if s in spec],
6118
+ [s for s in await tp.get_installable_versions(spec) if s in spec],
5259
6119
  key=lambda s: s.version,
5260
6120
  )
5261
6121
  if not sv:
5262
6122
  return None
5263
6123
 
5264
6124
  bv = sv[-1]
5265
- return tp.install_version(bv)
6125
+ return await tp.install_version(bv)
5266
6126
 
5267
- def list(self, spec: InterpSpecifier) -> None:
6127
+ async def list(self, spec: InterpSpecifier) -> None:
5268
6128
  print('installed:')
5269
6129
  for n, p in self._providers.items():
5270
6130
  lst = [
5271
6131
  si
5272
- for si in p.get_installed_versions(spec)
6132
+ for si in await p.get_installed_versions(spec)
5273
6133
  if spec.contains(si)
5274
6134
  ]
5275
6135
  if lst:
@@ -5283,7 +6143,7 @@ class InterpResolver:
5283
6143
  for n, p in self._providers.items():
5284
6144
  lst = [
5285
6145
  si
5286
- for si in p.get_installable_versions(spec)
6146
+ for si in await p.get_installable_versions(spec)
5287
6147
  if spec.contains(si)
5288
6148
  ]
5289
6149
  if lst:
@@ -5325,9 +6185,9 @@ class InterpCommand(Command['InterpCommand.Output']):
5325
6185
 
5326
6186
 
5327
6187
  class InterpCommandExecutor(CommandExecutor[InterpCommand, InterpCommand.Output]):
5328
- def execute(self, cmd: InterpCommand) -> InterpCommand.Output:
6188
+ async def execute(self, cmd: InterpCommand) -> InterpCommand.Output:
5329
6189
  i = InterpSpecifier.parse(check_not_none(cmd.spec))
5330
- o = check_not_none(DEFAULT_INTERP_RESOLVER.resolve(i, install=cmd.install))
6190
+ o = check_not_none(await DEFAULT_INTERP_RESOLVER.resolve(i, install=cmd.install))
5331
6191
  return InterpCommand.Output(
5332
6192
  exe=o.exe,
5333
6193
  version=str(o.version.version),
@@ -5366,7 +6226,7 @@ def bind_command(
5366
6226
  class _FactoryCommandExecutor(CommandExecutor):
5367
6227
  factory: ta.Callable[[], CommandExecutor]
5368
6228
 
5369
- def execute(self, i: Command) -> Command.Output:
6229
+ def execute(self, i: Command) -> ta.Awaitable[Command.Output]:
5370
6230
  return self.factory().execute(i)
5371
6231
 
5372
6232
 
@@ -5521,31 +6381,7 @@ def main_bootstrap(bs: MainBootstrap) -> Injector:
5521
6381
  ##
5522
6382
 
5523
6383
 
5524
- def _main() -> None:
5525
- import argparse
5526
-
5527
- parser = argparse.ArgumentParser()
5528
-
5529
- parser.add_argument('--_payload-file')
5530
-
5531
- parser.add_argument('-s', '--shell')
5532
- parser.add_argument('-q', '--shell-quote', action='store_true')
5533
- parser.add_argument('--python', default='python3')
5534
-
5535
- parser.add_argument('--pycharm-debug-port', type=int)
5536
- parser.add_argument('--pycharm-debug-host')
5537
- parser.add_argument('--pycharm-debug-version')
5538
-
5539
- parser.add_argument('--debug', action='store_true')
5540
-
5541
- parser.add_argument('--local', action='store_true')
5542
-
5543
- parser.add_argument('command', nargs='+')
5544
-
5545
- args = parser.parse_args()
5546
-
5547
- #
5548
-
6384
+ async def _async_main(args: ta.Any) -> None:
5549
6385
  bs = MainBootstrap(
5550
6386
  main_config=MainConfig(
5551
6387
  log_level='DEBUG' if args.debug else 'INFO',
@@ -5558,12 +6394,16 @@ def _main() -> None:
5558
6394
 
5559
6395
  pycharm_remote_debug=PycharmRemoteDebug(
5560
6396
  port=args.pycharm_debug_port,
5561
- host=args.pycharm_debug_host,
6397
+ **(dict(host=args.pycharm_debug_host) if args.pycharm_debug_host is not None else {}),
5562
6398
  install_version=args.pycharm_debug_version,
5563
6399
  ) if args.pycharm_debug_port is not None else None,
6400
+
6401
+ timebomb_delay_s=args.remote_timebomb_delay_s,
5564
6402
  ),
5565
6403
  )
5566
6404
 
6405
+ #
6406
+
5567
6407
  injector = main_bootstrap(
5568
6408
  bs,
5569
6409
  )
@@ -5582,7 +6422,7 @@ def _main() -> None:
5582
6422
 
5583
6423
  #
5584
6424
 
5585
- with contextlib.ExitStack() as es:
6425
+ async with contextlib.AsyncExitStack() as es:
5586
6426
  ce: CommandExecutor
5587
6427
 
5588
6428
  if args.local:
@@ -5595,16 +6435,51 @@ def _main() -> None:
5595
6435
  python=args.python,
5596
6436
  )
5597
6437
 
5598
- ce = es.enter_context(injector[RemoteExecution].connect(tgt, bs)) # noqa
6438
+ ce = await es.enter_async_context(injector[RemoteExecutionConnector].connect(tgt, bs)) # noqa
5599
6439
 
5600
- for cmd in cmds:
5601
- r = ce.try_execute(
6440
+ async def run_command(cmd: Command) -> None:
6441
+ res = await ce.try_execute(
5602
6442
  cmd,
5603
6443
  log=log,
5604
6444
  omit_exc_object=True,
5605
6445
  )
5606
6446
 
5607
- print(msh.marshal_obj(r, opts=ObjMarshalOptions(raw_bytes=True)))
6447
+ print(msh.marshal_obj(res, opts=ObjMarshalOptions(raw_bytes=True)))
6448
+
6449
+ await asyncio.gather(*[
6450
+ run_command(cmd)
6451
+ for cmd in cmds
6452
+ ])
6453
+
6454
+
6455
+ def _main() -> None:
6456
+ import argparse
6457
+
6458
+ parser = argparse.ArgumentParser()
6459
+
6460
+ parser.add_argument('--_payload-file')
6461
+
6462
+ parser.add_argument('-s', '--shell')
6463
+ parser.add_argument('-q', '--shell-quote', action='store_true')
6464
+ parser.add_argument('--python', default='python3')
6465
+
6466
+ parser.add_argument('--pycharm-debug-port', type=int)
6467
+ parser.add_argument('--pycharm-debug-host')
6468
+ parser.add_argument('--pycharm-debug-version')
6469
+
6470
+ parser.add_argument('--remote-timebomb-delay-s', type=float)
6471
+
6472
+ parser.add_argument('--debug', action='store_true')
6473
+
6474
+ parser.add_argument('--local', action='store_true')
6475
+
6476
+ parser.add_argument('command', nargs='+')
6477
+
6478
+ args = parser.parse_args()
6479
+
6480
+ #
6481
+
6482
+ asyncio.run(_async_main(args))
5608
6483
 
5609
6484
 
5610
6485
  if __name__ == '__main__':