chore: adjusts type hints to account for complications with mypy

Author: chalmerlowe

packages/bigframes/bigframes/core/compile/ibis_compiler/aggregate_compiler.py

@@ -528,8 +528,9 @@ def _(
     column: ibis_types.Column,
     window=None,
 ) -> ibis_types.Value:
+    # Ibis FirstNonNullValue expects Value[Any, Columnar], Mypy struggles to see Column as compatible.
     return _apply_window_if_present(
-        ibis_ops.FirstNonNullValue(column).to_expr(),
+        ibis_ops.FirstNonNullValue(column).to_expr(),  # type: ignore[arg-type]
         window,  # type: ignore
     )
 
@@ -549,8 +550,9 @@ def _(
     column: ibis_types.Column,
     window=None,
 ) -> ibis_types.Value:
+    # Ibis LastNonNullValue expects Value[Any, Columnar], Mypy struggles to see Column as compatible.
     return _apply_window_if_present(
-        ibis_ops.LastNonNullValue(column).to_expr(),
+        ibis_ops.LastNonNullValue(column).to_expr(),  # type: ignore[arg-type]
         window,  # type: ignore
     )
 
@@ -803,8 +805,9 @@ def _to_ibis_boundary(
 ) -> Optional[ibis_expr_window.WindowBoundary]:
     if boundary is None:
         return None
+    # WindowBoundary expects Value[Any, Any], ibis_types.literal returns Scalar which Mypy doesn't see as compatible.
     return ibis_expr_window.WindowBoundary(
-        abs(boundary),
+        ibis_types.literal(boundary if boundary >= 0 else -boundary),  # type: ignore[arg-type]
         preceding=boundary <= 0,  # type:ignore
     )
 

packages/bigframes/bigframes/core/compile/ibis_compiler/operations/geo_ops.py

@@ -182,7 +182,7 @@ def st_buffer(
 
 @ibis_udf.scalar.builtin
 def st_distance(
-    a: ibis_dtypes.geography, b: ibis_dtypes.geography, use_spheroid: bool
+    a: ibis_dtypes.geography, b: ibis_dtypes.geography, use_spheroid: bool  # type: ignore
 ) -> ibis_dtypes.float:  # type: ignore
     """Convert string to geography."""
 

packages/bigframes/bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -2168,9 +2168,12 @@ def obj_make_ref_json(objectref_json: ibis_dtypes.JSON) -> _OBJ_REF_IBIS_DTYPE:
 
 
 @ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL")
-def obj_get_access_url(
-    obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String
-) -> ibis_dtypes.JSON:  # type: ignore
+# Stub for BigQuery UDF, empty body is intentional.
+# _OBJ_REF_IBIS_DTYPE is a variable holding a type, Mypy complains about it being used as type hint.
+def obj_get_access_url(  # type: ignore[empty-body]
+    obj_ref: _OBJ_REF_IBIS_DTYPE,  # type: ignore[valid-type]
+    mode: ibis_dtypes.String
+) -> ibis_dtypes.JSON:
     """Get access url (as ObjectRefRumtime JSON) from ObjectRef."""
 
 

packages/bigframes/bigframes/core/expression_factoring.py

@@ -243,7 +243,7 @@ def factor_aggregation(root: nodes.ColumnDef) -> FactoredAggregation:
     }
 
     root_scalar_expr = nodes.ColumnDef(
-        sub_expressions(root.expression, agg_outputs_dict),
+        sub_expressions(root.expression, cast(Mapping[expression.Expression, expression.Expression], agg_outputs_dict)),
         root.id,  # type: ignore
     )
 

packages/bigframes/bigframes/core/local_data.py

@@ -33,6 +33,7 @@
 
 import bigframes.core.schema as schemata
 import bigframes.dtypes
+from bigframes.core import identifiers
 from bigframes.core import pyarrow_utils
 
 

packages/bigframes/bigframes/core/nodes.py

@@ -674,7 +674,7 @@ def fields(self) -> Sequence[Field]:
             Field(
                 col_id,
                 self.local_data_source.schema.get_type(source_id),
-                nullable=self.local_data_source.is_nullable(source_id),
+                nullable=self.local_data_source.is_nullable(identifiers.ColumnId(source_id)),
             )
             for col_id, source_id in self.scan_list.items
         )

packages/bigframes/bigframes/core/rewrite/as_sql.py

@@ -291,8 +291,9 @@ def _extract_ctes_to_with_expr(
         root.top_down(lambda x: mapping.get(x, x)),
         cte_names,
         tuple(
-            cte_node.child.top_down(lambda x: mapping.get(x, x))
-            for cte_node in topological_ctes  # type: ignore
+            # Mypy loses context that cte_node is a CteNode with a child attribute, despite the isinstance filter above.
+            cte_node.child.top_down(lambda x: mapping.get(x, x))  # type: ignore[attr-defined]
+            for cte_node in topological_ctes
         ),
     )
 

packages/bigframes/bigframes/session/iceberg.py

@@ -98,9 +98,10 @@ def _extract_location_from_catalog_extension_data(data):
 
 
 class SchemaVisitor(pyiceberg.schema.SchemaVisitorPerPrimitiveType[bq.SchemaField]):
-    def schema(
+    # Override returns a tuple of fields instead of a single field, violating supertype signature but intentional for this visitor.
+    def schema(  # type: ignore[override]
         self, schema: pyiceberg.schema.Schema, struct_result: bq.SchemaField
-    ) -> tuple[bq.SchemaField, ...]:  # type: ignore
+    ) -> tuple[bq.SchemaField, ...]:
         return tuple(f for f in struct_result.fields)
 
     def struct(