cleanup [3/n]: mypy-protobuf & minor implementation changes

ghstack-source-id: 735f901af37e65afb948b145e8ea4ff2d763114f
Pull Request resolved: #29

.pre-commit-config.yaml

@@ -22,12 +22,6 @@ repos:
          -   --disallow-incomplete-defs
          -   --ignore-missing-imports
          -   --no-warn-no-return
-         -   --disable-error-code
-         -   attr-defined
-         -   --disable-error-code
-         -   misc
-         -   --disable-error-code
-         -   var-annotated
          exclude: '^src/.*_pb2.py$'
          additional_dependencies: [types-protobuf]
 -   repo: https://github.com/codespell-project/codespell

src/isolate/backends/__init__.py

@@ -1,2 +1,2 @@
 from isolate.backends._base import *
-from isolate.backends.context import IsolateSettings
+from isolate.backends.settings import IsolateSettings

src/isolate/backends/_base.py

@@ -13,13 +13,8 @@
     TypeVar,
 )
 
-from isolate.backends.context import (
-    DEFAULT_SETTINGS,
-    IsolateSettings,
-    Log,
-    LogLevel,
-    LogSource,
-)
+from isolate.backends.settings import DEFAULT_SETTINGS, IsolateSettings
+from isolate.logs import Log, LogLevel, LogSource
 
 __all__ = [
     "BasicCallable",

src/isolate/backends/common.py

@@ -1,16 +1,13 @@
 from __future__ import annotations
 
-import functools
 import hashlib
-import importlib
 import os
 import shutil
-import sysconfig
 import threading
 from contextlib import contextmanager
 from functools import lru_cache
 from pathlib import Path
-from typing import Any, Callable, Iterator, Optional, Tuple
+from typing import Callable, Iterator, Optional, Tuple
 
 _OLD_DIR_PREFIX = "old-"
 
@@ -43,23 +40,6 @@ def replace_dir(src_path: Path, dst_path: Path) -> None:
             shutil.rmtree(cleanup)
 
 
-def python_path_for(*search_paths: Path) -> str:
-    """Return the PYTHONPATH for the library paths residing
-    in the given 'search_paths'. The order of the paths is
-    preserved."""
-
-    assert len(search_paths) >= 1
-    return os.pathsep.join(
-        # sysconfig defines the schema of the directories under
-        # any comforming Python installation (like venv, conda, etc.).
-        #
-        # Be aware that Debian's system installation does not
-        # comform sysconfig.
-        sysconfig.get_path("purelib", vars={"base": search_path})
-        for search_path in search_paths
-    )
-
-
 def get_executable_path(search_path: Path, executable_name: str) -> Path:
     """Return the path for the executable named 'executable_name' under
     the '/bin' directory of 'search_path'."""
@@ -75,23 +55,6 @@ def get_executable_path(search_path: Path, executable_name: str) -> Path:
     return Path(executable_path)
 
 
-_NO_HIT = object()
-
-
-def cache_static(func):
-    """Cache the result of an parameter-less function."""
-    _function_cache = _NO_HIT
-
-    @functools.wraps(func)
-    def wrapper():
-        nonlocal _function_cache
-        if _function_cache is _NO_HIT:
-            _function_cache = func()
-        return _function_cache
-
-    return wrapper
-
-
 _MESSAGE_STREAM_DELAY = 0.1
 
 
@@ -175,15 +138,6 @@ def logged_io(
             raise RuntimeError("Log observers did not terminate in time.")
 
 
-def run_serialized(serialization_method: str, data: bytes) -> Any:
-    """Deserialize the given 'data' into an parameter-less callable and
-    run it."""
-
-    serialization_backend = importlib.import_module(serialization_method)
-    executable = serialization_backend.loads(data)
-    return executable()
-
-
 @lru_cache(maxsize=None)
 def sha256_digest_of(*unique_fields: str, _join_char: str = "\n") -> str:
     """Return the SHA256 digest that corresponds to the combined version

src/isolate/backends/conda.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import os
 import shutil
 import subprocess
@@ -8,8 +9,8 @@
 from typing import Any, ClassVar, Dict, List
 
 from isolate.backends import BaseEnvironment, EnvironmentCreationError
-from isolate.backends.common import cache_static, logged_io, sha256_digest_of
-from isolate.backends.context import DEFAULT_SETTINGS, IsolateSettings
+from isolate.backends.common import logged_io, sha256_digest_of
+from isolate.backends.settings import DEFAULT_SETTINGS, IsolateSettings
 from isolate.connections import PythonIPC
 
 # Specify the path where the conda binary might reside in (or
@@ -86,7 +87,7 @@ def open_connection(self, connection_key: Path) -> PythonIPC:
         return PythonIPC(self, connection_key)
 
 
-@cache_static
+@functools.lru_cache(1)
 def _get_conda_executable() -> Path:
     for path in [_ISOLATE_CONDA_HOME, None]:
         conda_path = shutil.which(_CONDA_COMMAND, path=path)

src/isolate/backends/local.py

@@ -7,7 +7,7 @@
 
 from isolate.backends import BaseEnvironment
 from isolate.backends.common import sha256_digest_of
-from isolate.backends.context import DEFAULT_SETTINGS, IsolateSettings
+from isolate.backends.settings import DEFAULT_SETTINGS, IsolateSettings
 from isolate.connections import PythonIPC
 
 

src/isolate/backends/context.py → src/isolate/backends/settings.py

@@ -3,80 +3,21 @@
 import shutil
 import tempfile
 from contextlib import contextmanager
-from dataclasses import dataclass, field, replace
-from enum import Enum
+from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable, Dict, Iterator, NewType, Optional
+from typing import TYPE_CHECKING, Callable, Iterator
 
 from platformdirs import user_cache_dir
 
 from isolate.backends.common import replace_dir
+from isolate.logs import Log
 
 if TYPE_CHECKING:
     from isolate.backends import BaseEnvironment
 
 _SYSTEM_TEMP_DIR = Path(tempfile.gettempdir())
 
 
-class LogSource(str, Enum):
-    """Represents where the log orinates from."""
-
-    # During the environment creation process (e.g. if the environment
-    # is already created/cached, then no logs from this source will be
-    # emitted).
-    BUILDER = "builder"
-
-    # During the environment execution process (from the server<->agent
-    # communication, mostly for debugging purposes).
-    BRIDGE = "bridge"
-
-    # From the user script itself (e.g. a print() call in the given
-    # function). The stream will be attached as level (stdout or stderr)
-    USER = "user"
-
-
-class LogLevel(str, Enum):
-    """Represents the log level."""
-
-    TRACE = "trace"
-    DEBUG = "debug"
-    INFO = "info"
-    WARNING = "warning"
-    ERROR = "error"
-
-    # For user scripts
-    STDOUT = "stdout"
-    STDERR = "stderr"
-
-
-@dataclass
-class Log:
-    """A structured log message with an option source and level."""
-
-    message: str
-    source: LogSource
-    level: LogLevel = LogLevel.INFO
-    bound_env: Optional[BaseEnvironment] = field(default=None, repr=False)
-
-    def serialize(self) -> Dict[str, str]:
-        return {
-            "message": self.message,
-            "source": format(self.source),
-            "level": format(self.level),
-        }
-
-    def __str__(self) -> str:
-        parts = []
-        if self.bound_env:
-            parts.append(f"[{self.bound_env.key[:6]}]")
-        else:
-            parts.append("[global]")
-
-        parts.append(f"[{self.source}]".ljust(10))
-        parts.append(f"[{self.level}]".ljust(10))
-        return " ".join(parts) + self.message
-
-
 @dataclass(frozen=True)
 class IsolateSettings:
     cache_dir: Path = Path(user_cache_dir("isolate", "isolate"))

src/isolate/backends/virtualenv.py

@@ -13,7 +13,7 @@
     logged_io,
     sha256_digest_of,
 )
-from isolate.backends.context import DEFAULT_SETTINGS, IsolateSettings
+from isolate.backends.settings import DEFAULT_SETTINGS, IsolateSettings
 from isolate.connections import PythonIPC
 
 

src/isolate/connections/_local/_base.py

@@ -2,6 +2,7 @@
 
 import os
 import subprocess
+import sysconfig
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from functools import partial
@@ -17,19 +18,32 @@
     Union,
 )
 
-from isolate.backends.common import (
-    get_executable_path,
-    logged_io,
-    python_path_for,
-)
-from isolate.backends.context import LogLevel
+from isolate.backends.common import get_executable_path, logged_io
+from isolate.logs import LogLevel
 
 if TYPE_CHECKING:
     from isolate.backends import BaseEnvironment
 
 ConnectionType = TypeVar("ConnectionType")
 
 
+def python_path_for(*search_paths: Path) -> str:
+    """Return the PYTHONPATH for the library paths residing
+    in the given 'search_paths'. The order of the paths is
+    preserved."""
+
+    assert len(search_paths) >= 1
+    return os.pathsep.join(
+        # sysconfig defines the schema of the directories under
+        # any comforming Python installation (like venv, conda, etc.).
+        #
+        # Be aware that Debian's system installation does not
+        # comform sysconfig.
+        sysconfig.get_path("purelib", vars={"base": search_path})
+        for search_path in search_paths
+    )
+
+
 @dataclass
 class PythonExecutionBase(Generic[ConnectionType]):
     """A generic Python execution implementation that can trigger a new process

src/isolate/connections/grpc/_base.py

@@ -11,11 +11,11 @@
     CallResultType,
     EnvironmentConnection,
 )
-from isolate.backends.context import Log, LogLevel, LogSource
 from isolate.connections._local import PythonExecutionBase, agent_startup
 from isolate.connections.common import serialize_object
 from isolate.connections.grpc import agent, definitions
 from isolate.connections.grpc.interface import from_grpc, to_grpc
+from isolate.logs import Log, LogLevel, LogSource
 
 
 class AgentError(Exception):

src/isolate/connections/grpc/agent.py

@@ -6,14 +6,15 @@
 from argparse import ArgumentParser
 from concurrent import futures
 from dataclasses import dataclass
-from typing import Iterator
+from typing import Iterator, cast
 
 import grpc
 from grpc import ServicerContext, StatusCode
 
 from isolate.connections.common import SerializationError, serialize_object
 from isolate.connections.grpc import definitions
-from isolate.connections.grpc.interface import from_grpc
+from isolate.connections.grpc.interface import from_grpc, to_grpc
+from isolate.logs import Log, LogLevel, LogSource
 
 
 @dataclass
@@ -62,13 +63,13 @@ def Run(
         try:
             definition = serialize_object(request.method, result)
         except SerializationError:
-            yield from self.log(traceback.format_exc(), level=definitions.ERROR)
+            yield from self.log(traceback.format_exc(), level=LogLevel.ERROR)
             return self.abort_with_msg(
                 "The result of the input function could not be serialized.",
                 context,
             )
         except BaseException:
-            yield from self.log(traceback.format_exc(), level=definitions.ERROR)
+            yield from self.log(traceback.format_exc(), level=LogLevel.ERROR)
             return self.abort_with_msg(
                 "An unexpected error occurred while serializing the result.", context
             )
@@ -88,10 +89,11 @@ def Run(
     def log(
         self,
         message: str,
-        level: definitions.LogLevel = definitions.TRACE,
-        source: definitions.LogSource = definitions.BRIDGE,
+        level: LogLevel = LogLevel.TRACE,
+        source: LogSource = LogSource.BRIDGE,
     ) -> Iterator[definitions.PartialRunResult]:
-        log = definitions.Log(message=message, level=level, source=source)
+        log = to_grpc(Log(message, level=level, source=source))
+        log = cast(definitions.Log, log)
         yield definitions.PartialRunResult(result=None, is_complete=False, logs=[log])
 
     def abort_with_msg(

src/isolate/connections/grpc/definitions/agent_pb2.pyi

@@ -1,5 +1,7 @@
-from isolate.connections.grpc.definitions import common_pb2 as _common_pb2
-from google.protobuf import descriptor as _descriptor
-from typing import ClassVar as _ClassVar
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import google.protobuf.descriptor
 
-DESCRIPTOR: _descriptor.FileDescriptor
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor