diff --git a/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py b/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py index 787341263fde..903bf38ab279 100644 --- a/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py +++ b/sdks/python/apache_beam/examples/inference/pytorch_image_classification_with_side_inputs.py @@ -115,7 +115,7 @@ class PostProcessor(beam.DoFn): Return filename, prediction and the model id used to perform the prediction """ - def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]: + def process(self, element: tuple[str, PredictionResult]) -> Iterable[Any]: filename, prediction_result = element prediction = torch.argmax(prediction_result.inference, dim=0) yield filename, prediction, prediction_result.model_id diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py index f0a23ddce02a..1e90c2974778 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py @@ -496,7 +496,8 @@ def _emit_query_ranges( self, start_ts: Timestamp, end_ts: Timestamp, - watermark_estimator: _PollWatermarkEstimator) -> Iterable[_QueryRange]: + watermark_estimator: _PollWatermarkEstimator + ) -> Iterable[TimestampedValue[_QueryRange]]: """Compute and yield _QueryRange elements, advancing estimator state.""" ranges = compute_ranges(start_ts, end_ts, self._change_function) _LOGGER.info( @@ -524,7 +525,7 @@ def process( restriction_tracker=beam.DoFn.RestrictionParam(), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( _PollWatermarkEstimatorProvider()) - ) -> Iterable[_QueryRange]: + ) -> Iterable[TimestampedValue[_QueryRange]]: now = time.time() start_ts = watermark_estimator.poll_cursor() @@ -1031,7 +1032,10 @@ class _DecompressArrowBatchesFn(beam.DoFn): def __init__(self, change_timestamp_column: str = 'change_timestamp') -> None: self._change_timestamp_column = change_timestamp_column - def process(self, element: tuple[bytes, bytes]) -> Iterable[dict[str, Any]]: + def process( + self, + element: tuple[bytes, + bytes]) -> Iterable[TimestampedValue[dict[str, Any]]]: schema_bytes, batch_bytes = element schema = pyarrow.ipc.read_schema(pyarrow.py_buffer(schema_bytes)) batch = pyarrow.ipc.read_record_batch( diff --git a/sdks/python/pyproject.toml b/sdks/python/pyproject.toml index e769bf2d91e2..5901d817c9c7 100644 --- a/sdks/python/pyproject.toml +++ b/sdks/python/pyproject.toml @@ -212,5 +212,4 @@ invalid-argument = "ignore" invalid-inheritance = "ignore" not-iterable = "ignore" unexpected-keyword = "ignore" -invalid-yield = "ignore" bad-typed-dict-key = "ignore"