ref: remove proto re-exports from types.py, import directly from tusk.drift.core.v1

types.py now only contains SDK wrapper dataclasses and utility functions.
Proto types are imported directly where used. Removes SDKMessageType/CLIMessageType aliases.

Author: sohil-kshirsagar

drift/core/communication/__init__.py

@@ -6,38 +6,25 @@
 
 from .communicator import CommunicatorConfig, ProtobufCommunicator
 from .types import (
-    CliMessage,
-    CLIMessageType,
     ConnectRequest,
     ConnectResponse,
     GetMockRequest,
     GetMockResponse,
-    MessageType,
     MockRequestInput,
     MockResponseOutput,
-    # Protobuf types (re-exported)
-    SdkMessage,
-    SDKMessageType,
     dict_to_span,
     extract_response_data,
     span_to_proto,
 )
 
 __all__ = [
-    # Message types
-    "MessageType",
-    "SDKMessageType",
-    "CLIMessageType",
     # Request/Response types
     "ConnectRequest",
     "ConnectResponse",
     "GetMockRequest",
     "GetMockResponse",
     "MockRequestInput",
     "MockResponseOutput",
-    # Protobuf types
-    "SdkMessage",
-    "CliMessage",
     # Utilities
     "span_to_proto",
     "dict_to_span",

drift/core/communication/communicator.py

@@ -10,27 +10,31 @@
 from dataclasses import dataclass
 from typing import Any
 
-from tusk.drift.core.v1 import GetMockRequest as ProtoGetMockRequest
-
-from ...version import MIN_CLI_VERSION, SDK_VERSION
-from ..span_serialization import clean_span_to_proto
-from ..types import CleanSpanData, calling_library_context
-from .types import (
+from tusk.drift.core.v1 import (
     BranchInfo,
     CliMessage,
-    ConnectRequest,
     CoverageSnapshotResponse,
     FileCoverageData,
-    GetMockRequest,
     InstrumentationVersionMismatchAlert,
     MessageType,
-    MockRequestInput,
-    MockResponseOutput,
     SdkMessage,
     SendAlertRequest,
     SendInboundSpanForReplayRequest,
     SetTimeTravelResponse,
     UnpatchedDependencyAlert,
+)
+from tusk.drift.core.v1 import (
+    GetMockRequest as ProtoGetMockRequest,
+)
+
+from ...version import MIN_CLI_VERSION, SDK_VERSION
+from ..span_serialization import clean_span_to_proto
+from ..types import CleanSpanData, calling_library_context
+from .types import (
+    ConnectRequest,
+    GetMockRequest,
+    MockRequestInput,
+    MockResponseOutput,
     span_to_proto,
 )
 

drift/core/communication/types.py

@@ -12,20 +12,6 @@
 from __future__ import annotations
 
 __all__ = [
-    # Re-exported protobuf types
-    "CliMessage",
-    "InstrumentationVersionMismatchAlert",
-    "MessageType",
-    "Runtime",
-    "SdkMessage",
-    "SendAlertRequest",
-    "SendInboundSpanForReplayRequest",
-    "SetTimeTravelRequest",
-    "SetTimeTravelResponse",
-    "UnpatchedDependencyAlert",
-    # Aliases
-    "SDKMessageType",
-    "CLIMessageType",
     # Dataclasses
     "ConnectRequest",
     "ConnectResponse",
@@ -42,21 +28,6 @@
 from dataclasses import dataclass, field
 from typing import Any
 
-from tusk.drift.core.v1 import (
-    BranchInfo,  # noqa: F401 (re-exported)
-    CliMessage,
-    CoverageSnapshotResponse,  # noqa: F401 (re-exported)
-    FileCoverageData,  # noqa: F401 (re-exported)
-    InstrumentationVersionMismatchAlert,
-    MessageType,
-    Runtime,
-    SdkMessage,
-    SendAlertRequest,
-    SendInboundSpanForReplayRequest,
-    SetTimeTravelRequest,
-    SetTimeTravelResponse,
-    UnpatchedDependencyAlert,
-)
 from tusk.drift.core.v1 import (
     ConnectRequest as ProtoConnectRequest,
 )
@@ -69,6 +40,9 @@
 from tusk.drift.core.v1 import (
     GetMockResponse as ProtoGetMockResponse,
 )
+from tusk.drift.core.v1 import (
+    Runtime,
+)
 from tusk.drift.core.v1 import (
     Span as ProtoSpan,
 )
@@ -82,9 +56,6 @@
     StatusCode as ProtoStatusCode,
 )
 
-SDKMessageType = MessageType
-CLIMessageType = MessageType
-
 
 def _python_to_value(value: Any) -> Any:
     """Convert Python value to protobuf Value."""