googleapis · tswast · May 15, 2024 · May 15, 2024 · May 15, 2024 · May 15, 2024
@@ -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,
         )

@@ -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]
 

@@ -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

@@ -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

@@ -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(

@@ -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()

@@ -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

@@ -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):

@@ -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
                 )
             )

@@ -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"',
     )

@@ -1040,7 +1040,7 @@ def floats_pd():
         dtype=pd.Float64Dtype(),
     )
     # Index helps debug failed cases
-    df.index = df.float64_col
+    df.index = df.float64_col  # type: ignore
     # Upload fails if index name same as column name
     df.index.name = None
     return df.float64_col
@@ -1050,7 +1050,7 @@ def floats_pd():
 def floats_product_pd(floats_pd):
     df = pd.merge(floats_pd, floats_pd, how="cross")
     # Index helps debug failed cases
-    df = df.set_index([df.float64_col_x, df.float64_col_y])
+    df = df.set_index([df.float64_col_x, df.float64_col_y])  # type: ignore
     df.index.names = ["left", "right"]
     return df
 

@@ -113,9 +113,10 @@ def test_cluster_configure_fit_score_predict(
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_cluster_model", replace=True
     )
+    # TODO(b/340875247): fix type error
     assert (
         f"{dataset_id}.temp_configured_cluster_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.n_clusters == 3
     assert reloaded_model.init == "RANDOM"
@@ -153,9 +154,10 @@ def test_cluster_configure_fit_load_params(penguins_df_default_index, dataset_id
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_cluster_model", replace=True
     )
+    # TODO(b/340875247): fix type error
     assert (
         f"{dataset_id}.temp_configured_cluster_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.n_clusters == 4
     assert reloaded_model.init == "RANDOM"

@@ -147,7 +147,8 @@ def test_columntransformer_save_load(new_penguins_df, dataset_id):
         new_penguins_df[["species", "culmen_length_mm", "flipper_length_mm"]]
     ).to_pandas()
 
-    expected = pandas.DataFrame(
+    # TODO(b/340888429): fix type error
+    expected = pandas.DataFrame(  # type: ignore
         {
             "onehotencoded_species": [
                 [{"index": 1, "value": 1.0}],
@@ -164,4 +165,5 @@ def test_columntransformer_save_load(new_penguins_df, dataset_id):
         index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"),
     )
 
-    pandas.testing.assert_frame_equal(result, expected, rtol=0.1, check_dtype=False)
+    # TODO(b/340888429): fix type error
+    pandas.testing.assert_frame_equal(result, expected, rtol=0.1, check_dtype=False)  # type: ignore
@@ -79,9 +79,10 @@ def test_decomposition_configure_fit_score_predict(
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_pca_model", replace=True
     )
+    # TODO(b/340876028): fix type error
     assert (
         f"{dataset_id}.temp_configured_pca_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.n_components == 3
 
@@ -149,9 +150,10 @@ def test_decomposition_configure_fit_score_predict_params(
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_pca_model", replace=True
     )
+    # TODO(b/340876028): fix type error
     assert (
         f"{dataset_id}.temp_configured_pca_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.n_components == 5
     assert reloaded_model.svd_solver == "RANDOMIZED"
@@ -167,9 +169,10 @@ def test_decomposition_configure_fit_load_float_component(
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_pca_model", replace=True
     )
+    # TODO(b/340876028): fix type error
     assert (
         f"{dataset_id}.temp_configured_pca_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.n_components == 0.2
 
@@ -184,8 +187,9 @@ def test_decomposition_configure_fit_load_none_component(
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_pca_model", replace=True
     )
+    # TODO(b/340876028): fix type error
     assert (
         f"{dataset_id}.temp_configured_pca_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.n_components == 7
@@ -60,9 +60,10 @@ def test_xgbregressor_default_params(penguins_df_default_index, dataset_id):
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_xgbregressor_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_xgbregressor_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
 
 
@@ -118,9 +119,10 @@ def test_xgbregressor_dart_booster_multiple_params(
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_xgbregressor_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_xgbregressor_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.booster == "DART"
     assert reloaded_model.dart_normalized_type == "TREE"
@@ -174,9 +176,10 @@ def test_xgbclassifier_default_params(penguins_df_default_index, dataset_id):
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_xgbclassifier_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_xgbclassifier_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
 
 
@@ -231,9 +234,10 @@ def test_xgbclassifier_dart_booster_multiple_params(
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_xgbclassifier_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_xgbclassifier_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.booster == "DART"
     assert reloaded_model.dart_normalized_type == "TREE"
@@ -288,9 +292,10 @@ def test_randomforestregressor_default_params(penguins_df_default_index, dataset
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_randomforestregressor_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_randomforestregressor_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
 
 
@@ -341,9 +346,10 @@ def test_randomforestregressor_multiple_params(penguins_df_default_index, datase
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_randomforestregressor_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_randomforestregressor_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.tree_method == "AUTO"
     assert reloaded_model.colsample_bytree == 0.95
@@ -394,16 +400,18 @@ def test_randomforestclassifier_default_params(penguins_df_default_index, datase
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_randomforestclassifier_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_randomforestclassifier_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
 
 
 @pytest.mark.flaky(retries=2)
 def test_randomforestclassifier_multiple_params(penguins_df_default_index, dataset_id):
+    # TODO(b/340888645): fix type error
     model = bigframes.ml.ensemble.RandomForestClassifier(
-        tree_method="AUTO",
+        tree_method="AUTO",  # type: ignore
         min_tree_child_weight=2,
         colsample_bytree=0.95,
         colsample_bylevel=0.95,
@@ -446,9 +454,10 @@ def test_randomforestclassifier_multiple_params(penguins_df_default_index, datas
     reloaded_model = model.to_gbq(
         f"{dataset_id}.temp_configured_randomforestclassifier_model", replace=True
     )
+    # TODO(b/340888645): fix type error
     assert (
         f"{dataset_id}.temp_configured_randomforestclassifier_model"
-        in reloaded_model._bqml_model.model_name
+        in reloaded_model._bqml_model.model_name  # type: ignore
     )
     assert reloaded_model.tree_method == "AUTO"
     assert reloaded_model.colsample_bytree == 0.95