Add mypy to the commit hooks and address typing issues (#72)

* Add mypy to the commit hooks and address typing issues

* Ignore `DataFrameType` and `SeriesType`

Author: karlhigley

.pre-commit-config.yaml

@@ -10,6 +10,14 @@ repos:
     - id: isort
       additional_dependencies: [toml]
       exclude: examples/*
+  # types
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: 'v0.940'
+    hooks:
+    - id: mypy
+      language_version: python3
+      args: [--no-strict-optional, --ignore-missing-imports, --show-traceback, --install-types, --non-interactive]
+      exclude: docs/*
   # code style
   - repo: https://github.com/python/black
     rev: 22.3.0

merlin/core/dispatch.py

@@ -72,11 +72,12 @@ def inner2(*args, **kwargs):
 
 
 if HAS_GPU:
-    DataFrameType = Union[pd.DataFrame, cudf.DataFrame]
-    SeriesType = Union[pd.Series, cudf.Series]
+    DataFrameType = Union[pd.DataFrame, cudf.DataFrame]  # type: ignore
+    SeriesType = Union[pd.Series, cudf.Series]  # type: ignore
 else:
-    DataFrameType = Union[pd.DataFrame]
-    SeriesType = Union[pd.Series]
+    DataFrameType = pd.DataFrame  # type: ignore
+    SeriesType = pd.Series  # type: ignore
+
 
 # Define mapping between non-nullable,
 # and nullable types in Pandas

merlin/io/csv.py

@@ -60,7 +60,7 @@ def to_ddf(self, columns=None, cpu=None):
             ddf = ddf[columns]
         return ddf
 
-    @property
+    @property  # type: ignore
     @functools.lru_cache(1)
     def _file_partition_map(self):
         ind = 0

merlin/io/parquet.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 #
 import functools
+import io as py_io
 import itertools
 import logging
 import math
@@ -22,7 +23,6 @@
 import threading
 import warnings
 from collections import defaultdict
-from io import BytesIO
 from uuid import uuid4
 
 from packaging.version import Version
@@ -332,12 +332,12 @@ def __init__(
 
         assert self.row_groups_per_part > 0
 
-    @property
+    @property  # type: ignore
     @functools.lru_cache(1)
     def _path0(self):
         return next(self._dataset.get_fragments()).path
 
-    @property
+    @property  # type: ignore
     @functools.lru_cache(1)
     def _legacy_dataset(self):
         # TODO: Remove this after finding a way to avoid
@@ -355,7 +355,7 @@ def _legacy_dataset(self):
             dataset = pq.ParquetDataset(paths[0], filesystem=fs)
         return dataset
 
-    @property
+    @property  # type: ignore
     @functools.lru_cache(1)
     def _dataset(self):
         paths = self.stripped_paths
@@ -595,7 +595,7 @@ def validate_dataset(
                     paths,
                     out_dir=False,
                 )
-                with BytesIO() as myio:
+                with py_io.BytesIO() as myio:
                     myio.write(memoryview(metadata_bytes))
                     myio.seek(0)
                     metadata = pq.ParquetFile(myio).metadata
@@ -930,9 +930,7 @@ def _write_data(data_list, output_path, fs, fn):
 class BaseParquetWriter(ThreadedWriter):
     def __init__(self, out_dir, suffix=".parquet", **kwargs):
         super().__init__(out_dir, **kwargs)
-        self.data_paths = []
         self.data_files = []
-        self.data_writers = []
         self.data_bios = []
         self._lock = threading.RLock()
         self.pwriter = self._pwriter
@@ -968,7 +966,7 @@ def _append_writer(self, path, schema=None, add_args=None, add_kwargs=None):
         _kwargs = tlz.merge(self.pwriter_kwargs, add_kwargs or {})
 
         if self.bytes_io:
-            bio = BytesIO()
+            bio = py_io.BytesIO()
             self.data_bios.append(bio)
             self.data_writers.append(self.pwriter(bio, *_args, **_kwargs))
         else:

merlin/io/writer.py

@@ -89,7 +89,9 @@ def __init__(
             self.num_out_files = num_out_files
         self.num_samples = [0] * self.num_out_files
 
-        self.data_paths = None
+        self.data_paths = []
+        self.data_writers = []
+
         self.need_cal_col_names = True
         self.use_guid = use_guid
         self.bytes_io = bytes_io

merlin/schema/io/tensorflow_metadata.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import os
 import pathlib
 from typing import Union
 
@@ -63,7 +64,7 @@ def from_json(cls, json: Union[str, bytes]) -> "TensorflowMetadata":
         return TensorflowMetadata(schema)
 
     @classmethod
-    def from_json_file(cls, path: str) -> "TensorflowMetadata":
+    def from_json_file(cls, path: os.PathLike) -> "TensorflowMetadata":
         """Create a TensorflowMetadata schema object from a JSON file
 
         Parameters
@@ -103,7 +104,9 @@ def from_proto_text(cls, proto_text: str) -> "TensorflowMetadata":
         return TensorflowMetadata(schema)
 
     @classmethod
-    def from_proto_text_file(cls, path: str, file_name="schema.pbtxt") -> "TensorflowMetadata":
+    def from_proto_text_file(
+        cls, path: os.PathLike, file_name="schema.pbtxt"
+    ) -> "TensorflowMetadata":
         """Create a TensorflowMetadata schema object from a Protobuf text file
 
         Parameters
@@ -120,7 +123,7 @@ def from_proto_text_file(cls, path: str, file_name="schema.pbtxt") -> "Tensorflo
 
         """
         path = pathlib.Path(path) / file_name
-        return cls.from_proto_text(_read_file(str(path)))
+        return cls.from_proto_text(_read_file(path))
 
     def to_proto_text(self) -> str:
         """Convert this TensorflowMetadata schema object to a Proto text string
@@ -414,7 +417,7 @@ def _merlin_column(feature):
     return ColumnSchema(name, tags, properties, dtype, is_list, is_ragged=is_ragged)
 
 
-def _read_file(path: str):
+def _read_file(path: os.PathLike):
     # TODO: Should we be using fsspec here too?
     path = pathlib.Path(path)
     if path.is_file():

merlin/schema/schema.py

@@ -35,7 +35,7 @@ class ColumnSchema:
 
     name: Text
     tags: Optional[TagSet] = field(default_factory=TagSet)
-    properties: Optional[Dict[str, any]] = field(default_factory=dict)
+    properties: Optional[Dict] = field(default_factory=dict)
     dtype: Optional[object] = None
     is_list: bool = False
     is_ragged: bool = False

merlin/schema/tags.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 #
 from enum import Enum
-from typing import List, Union
+from typing import List, Set, Union
 
 
 class Tags(Enum):
@@ -55,11 +55,12 @@ class TagSet:
     """Collection that normalizes tags and prevents collisions between incompatible tags"""
 
     def __init__(self, tags: List[Union[str, Tags]] = None):
-        tags = tags or []
         if isinstance(tags, TagSet):
-            tags = tags._tags
+            tags = list(tags._tags)
+        elif tags is None:
+            tags = []
 
-        self._tags = self._normalize_tags(tags)
+        self._tags: Set[Tags] = self._normalize_tags(tags)
 
         collisions = self._detect_collisions(self._tags, self._tags)
         if collisions:
@@ -121,7 +122,7 @@ def _convert_to_tagset(self, tags):
 
         return tags
 
-    def _normalize_tags(self, tags):
+    def _normalize_tags(self, tags) -> Set[Tags]:
         return set(Tags[tag.upper()] if tag in Tags._value2member_map_ else tag for tag in tags)
 
     def __repr__(self) -> str: