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,12 +43,12 @@ 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.

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

@@ -116,7 +116,7 @@ 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