Fix type annotation of protos.

.github/workflows/python-ci.yml

@@ -36,7 +36,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install mypy pylint pytest mock
-    - name: Install runtime dependencies
+    - name: Install runtime dependencies and Space
       working-directory: ./python
       run: |
         pip install .[dev]

python/src/space/core/fs/arrow.py

@@ -16,11 +16,10 @@
 
 from abc import abstractmethod
 
-from google.protobuf import message
 from google.protobuf import text_format
 from pyarrow import fs
 
-from space.core.fs.base import BaseFileSystem
+from space.core.fs.base import BaseFileSystem, ProtoT
 from space.core.utils.protos import proto_to_text
 from space.core.utils.uuids import random_id
 
@@ -39,7 +38,7 @@ def create_fs(self) -> fs.FileSystem:
   def create_dir(self, dir_path: str) -> None:
     self._fs.create_dir(dir_path)
 
-  def write_proto(self, file_path: str, msg: message.Message) -> None:
+  def write_proto(self, file_path: str, msg: ProtoT) -> None:
     # TODO: the current implement overwrite an existing file; to support an
     # to disallow overwrite.
     tmp_file_path = f"{file_path}.{random_id()}.tmp"
@@ -49,8 +48,7 @@ def write_proto(self, file_path: str, msg: message.Message) -> None:
 
     self._fs.move(tmp_file_path, file_path)
 
-  def read_proto(self, file_path: str,
-                 empty_msg: message.Message) -> message.Message:
+  def read_proto(self, file_path: str, empty_msg: ProtoT) -> ProtoT:
     with self._fs.open_input_file(file_path) as f:
       result = text_format.Parse(f.readall(), empty_msg)
     return result

python/src/space/core/fs/base.py

@@ -15,9 +15,12 @@
 """Abstract base file system."""
 
 from abc import ABC, abstractmethod
+from typing import TypeVar
 
 from google.protobuf import message
 
+ProtoT = TypeVar("ProtoT", bound=message.Message)
+
 
 class BaseFileSystem(ABC):
   """Abstract file system."""
@@ -27,10 +30,9 @@ def create_dir(self, dir_path: str) -> None:
     """Create a new directory."""
 
   @abstractmethod
-  def write_proto(self, file_path: str, msg: message.Message) -> None:
+  def write_proto(self, file_path: str, msg: ProtoT) -> None:
     """Write a proto message in text format to a file."""
 
   @abstractmethod
-  def read_proto(self, file_path: str,
-                 empty_msg: message.Message) -> message.Message:
+  def read_proto(self, file_path: str, empty_msg: ProtoT) -> ProtoT:
     """Read a proto message in text format from a file."""

python/src/space/core/proto/metadata.proto

@@ -43,18 +43,18 @@ message StorageMetadata {
   // The storage type.
   enum Type {
     TYPE_UNSPECIFIED = 0;
-    // A Space dataset.
+    // The dataset type supports fully managed storage features.
     DATASET = 1;
   }
   Type type = 3;
 
-  // The dataset schema.
+  // The storage schema.
   Schema schema = 4;
 
   // The current snapshot ID.
   int64 current_snapshot_id = 5;
 
-  // All alive snapshots.
+  // All alive snapshots with snapshot ID as key.
   map<int64, Snapshot> snapshots = 6;
 }
 

python/src/space/core/proto/metadata_pb2.pyi

@@ -74,14 +74,14 @@ class StorageMetadata(google.protobuf.message.Message):
         DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
         TYPE_UNSPECIFIED: StorageMetadata._Type.ValueType  # 0
         DATASET: StorageMetadata._Type.ValueType  # 1
-        """A Space dataset."""
+        """The dataset type supports fully managed storage features."""
 
     class Type(_Type, metaclass=_TypeEnumTypeWrapper):
         """The storage type."""
 
     TYPE_UNSPECIFIED: StorageMetadata.Type.ValueType  # 0
     DATASET: StorageMetadata.Type.ValueType  # 1
-    """A Space dataset."""
+    """The dataset type supports fully managed storage features."""
 
     @typing_extensions.final
     class SnapshotsEntry(google.protobuf.message.Message):
@@ -116,12 +116,12 @@ class StorageMetadata(google.protobuf.message.Message):
     type: global___StorageMetadata.Type.ValueType
     @property
     def schema(self) -> global___Schema:
-        """The dataset schema."""
+        """The storage schema."""
     current_snapshot_id: builtins.int
     """The current snapshot ID."""
     @property
     def snapshots(self) -> google.protobuf.internal.containers.MessageMap[builtins.int, global___Snapshot]:
-        """All alive snapshots."""
+        """All alive snapshots with snapshot ID as key."""
     def __init__(
         self,
         *,

python/src/space/core/storage.py

@@ -60,7 +60,9 @@ def metadata(self) -> meta.StorageMetadata:
     return self._metadata
 
   def snapshot(self, snapshot_id: Optional[int] = None) -> meta.Snapshot:
-    """Return the snapshot specified by a snapshot ID."""
+    """Return the snapshot specified by a snapshot ID, or current snapshot ID
+    if not specified.
+    """
     if snapshot_id is None:
       snapshot_id = self._metadata.current_snapshot_id