chore: type check with mypy --check-untyped-defs option (#693)

* chore: type check with mypy `--check-untyped-defs` option

* add TODOs for tests/system/small/ml/test_metrics.py

* add TODOs for tests/system/large/ml/test_cluster.py

* add TODOs for tests/system/large/ml/test_decomposition.py

* more TODOs

* test_dataframe TODOs

* tests/system/small/test_encryption.py TODOs

* test_index TODOs

* TODO test_index, test_multiindex

* tests/system/small/test_pandas.py TODOs

* test_series TODOs

* more TODOs

* tests/system/large/ml/test_ensemble.py TODOs

* more TODOs

* more TODOs

* add final TODOs

* add bug number

Author: tswast

bigframes/core/blocks.py

@@ -1226,8 +1226,9 @@ def calculate_pairwise_metric(self, op=agg_ops.CorrOp()):
         )
         labels = self._get_labels_for_columns(self.value_columns)
 
+        # TODO(b/340896143): fix type error
         expr = expr.unpivot(
-            row_labels=labels,
+            row_labels=labels,  # type: ignore
             index_col_ids=index_col_ids,
             unpivot_columns=unpivot_columns,
         )

bigframes/core/indexers.py

@@ -160,7 +160,8 @@ def __getitem__(self, key):
 
                 columns = key[1]
                 if isinstance(columns, pd.Series) and columns.dtype == "bool":
-                    columns = df.columns[columns]
+                    # TODO(b/340892590): fix type error
+                    columns = df.columns[columns]  # type: ignore
 
                 return df[columns]
 

bigframes/core/indexes/base.py

@@ -91,7 +91,8 @@ def __new__(
         from bigframes.core.indexes.multi import MultiIndex
 
         klass = MultiIndex if len(block._index_columns) > 1 else cls
-        result = typing.cast(Index, object.__new__(klass))
+        # TODO(b/340893286): fix type error
+        result = typing.cast(Index, object.__new__(klass))  # type: ignore
         result._query_job = None
         result._block = block
         return result

bigframes/dataframe.py

@@ -1762,7 +1762,8 @@ def label_filter(label):
                 if like:
                     return like in label_str
                 else:  # regex
-                    return re.match(regex, label_str) is not None
+                    # TODO(b/340891296): fix type error
+                    return re.match(regex, label_str) is not None  # type: ignore
 
             cols = [
                 col_id

bigframes/functions/remote_function.py

@@ -1023,8 +1023,9 @@ def wrapper(f):
             raise TypeError("f must be callable, got {}".format(f))
 
         signature = inspect.signature(f)
+        # TODO(b/340898611): fix type error
         ibis_signature = ibis_signature_from_python_signature(
-            signature, input_types, output_type
+            signature, input_types, output_type  # type: ignore
         )
 
         remote_function_client = RemoteFunctionClient(

bigframes/operations/__init__.py

@@ -491,7 +491,8 @@ def output_type(self, *input_types):
         if self.to_type == pa.string():
             return dtypes.STRING_DTYPE
         if isinstance(self.to_type, str):
-            return dtypes.BIGFRAMES_STRING_TO_BIGFRAMES[self.to_type]
+            # TODO(b/340895446): fix type error
+            return dtypes.BIGFRAMES_STRING_TO_BIGFRAMES[self.to_type]  # type: ignore
         return self.to_type
 
 
@@ -513,7 +514,8 @@ class RemoteFunctionOp(UnaryOp):
 
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
-        return self.func.output_dtype
+        # TODO(b/340895446): fix type error
+        return self.func.output_dtype  # type: ignore
 
 
 @dataclasses.dataclass(frozen=True)
@@ -627,7 +629,8 @@ class BinaryRemoteFunctionOp(BinaryOp):
 
     def output_type(self, *input_types):
         # This property should be set to a valid Dtype by the @remote_function decorator or read_gbq_function method
-        return self.func.output_dtype
+        # TODO(b/340895446): fix type error
+        return self.func.output_dtype  # type: ignore
 
 
 add_op = AddOp()

bigframes/operations/_matplotlib/__init__.py

@@ -24,7 +24,8 @@
 
 
 def plot(data, kind, **kwargs):
-    plot_obj = PLOT_CLASSES[kind](data, **kwargs)
+    # TODO(b/340896123): fix type error
+    plot_obj = PLOT_CLASSES[kind](data, **kwargs)  # type: ignore
     plot_obj.generate()
     plot_obj.draw()
     return plot_obj.result

bigframes/operations/_matplotlib/core.py

@@ -39,7 +39,8 @@ def draw(self) -> None:
 
     @property
     def result(self):
-        return self.axes
+        # TODO(b/340896123): fix type error
+        return self.axes  # type: ignore
 
 
 class SamplingPlot(MPLPlot):

bigframes/session/clients.py

@@ -117,10 +117,11 @@ def __init__(
     def _create_bigquery_client(self):
         bq_options = None
         if self._use_regional_endpoints:
+            # TODO(b/340896138): fix type error
             bq_options = google.api_core.client_options.ClientOptions(
                 api_endpoint=(
                     _BIGQUERY_REGIONAL_ENDPOINT
-                    if self._location.lower() in _REP_SUPPORTED_REGIONS
+                    if self._location.lower() in _REP_SUPPORTED_REGIONS  # type: ignore
                     else _BIGQUERY_LOCATIONAL_ENDPOINT
                 ).format(location=self._location),
             )
@@ -158,12 +159,11 @@ def bqconnectionclient(self):
             bqconnection_info = google.api_core.gapic_v1.client_info.ClientInfo(
                 user_agent=self._application_name
             )
-            self._bqconnectionclient = (
-                google.cloud.bigquery_connection_v1.ConnectionServiceClient(
-                    client_info=bqconnection_info,
-                    client_options=bqconnection_options,
-                    credentials=self._credentials,
-                )
+            # TODO(b/340896138): fix type error
+            self._bqconnectionclient = google.cloud.bigquery_connection_v1.ConnectionServiceClient(  # type: ignore
+                client_info=bqconnection_info,
+                client_options=bqconnection_options,
+                credentials=self._credentials,
             )
 
         return self._bqconnectionclient
@@ -173,18 +173,20 @@ def bqstoragereadclient(self):
         if not self._bqstoragereadclient:
             bqstorage_options = None
             if self._use_regional_endpoints:
+                # TODO(b/340896138): fix type error
                 bqstorage_options = google.api_core.client_options.ClientOptions(
                     api_endpoint=(
                         _BIGQUERYSTORAGE_REGIONAL_ENDPOINT
-                        if self._location.lower() in _REP_SUPPORTED_REGIONS
+                        if self._location.lower() in _REP_SUPPORTED_REGIONS  # type: ignore
                         else _BIGQUERYSTORAGE_LOCATIONAL_ENDPOINT
                     ).format(location=self._location),
                 )
             bqstorage_info = google.api_core.gapic_v1.client_info.ClientInfo(
                 user_agent=self._application_name
             )
+            # TODO(b/340896138): fix type error
             self._bqstoragereadclient = (
-                google.cloud.bigquery_storage_v1.BigQueryReadClient(
+                google.cloud.bigquery_storage_v1.BigQueryReadClient(  # type: ignore
                     client_info=bqstorage_info,
                     client_options=bqstorage_options,
                     credentials=self._credentials,
@@ -199,8 +201,9 @@ def cloudfunctionsclient(self):
             functions_info = google.api_core.gapic_v1.client_info.ClientInfo(
                 user_agent=self._application_name
             )
+            # TODO(b/340896138): fix type error
             self._cloudfunctionsclient = (
-                google.cloud.functions_v2.FunctionServiceClient(
+                google.cloud.functions_v2.FunctionServiceClient(  # type: ignore
                     client_info=functions_info,
                     credentials=self._credentials,
                 )
@@ -214,8 +217,9 @@ def resourcemanagerclient(self):
             resourcemanager_info = google.api_core.gapic_v1.client_info.ClientInfo(
                 user_agent=self._application_name
             )
+            # TODO(b/340896138): fix type error
             self._resourcemanagerclient = (
-                google.cloud.resourcemanager_v3.ProjectsClient(
+                google.cloud.resourcemanager_v3.ProjectsClient(  # type: ignore
                     credentials=self._credentials, client_info=resourcemanager_info
                 )
             )

noxfile.py

@@ -248,6 +248,7 @@ def mypy(session):
         "bigframes",
         os.path.join("tests", "system"),
         os.path.join("tests", "unit"),
+        "--check-untyped-defs",
         "--explicit-package-bases",
         '--exclude="^third_party"',
     )