You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by jr...@apache.org on 2025/01/09 16:02:58 UTC
(beam) 01/01: Finish yapf upgrade to 0.43.0
This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch yapfUpgrade
in repository https://gitbox.apache.org/repos/asf/beam.git
commit f4dc6c657519ee16ed323b1521eee27dfbafe36e
Author: Jack McCluskey <th...@gmail.com>
AuthorDate: Thu Jan 9 11:02:34 2025 -0500
Finish yapf upgrade to 0.43.0
---
sdks/python/apache_beam/coders/avro_record.py | 1 +
sdks/python/apache_beam/coders/coder_impl.py | 48 ++++-
sdks/python/apache_beam/coders/coders.py | 49 ++++-
.../coders/coders_property_based_test.py | 6 +-
sdks/python/apache_beam/coders/coders_test.py | 11 +
.../apache_beam/coders/coders_test_common.py | 17 +-
sdks/python/apache_beam/coders/fast_coders_test.py | 1 +
sdks/python/apache_beam/coders/observable.py | 1 +
sdks/python/apache_beam/coders/observable_test.py | 2 +
sdks/python/apache_beam/coders/row_coder.py | 2 +
sdks/python/apache_beam/coders/row_coder_test.py | 3 +-
sdks/python/apache_beam/coders/slow_coders_test.py | 1 +
sdks/python/apache_beam/coders/slow_stream.py | 3 +
.../apache_beam/coders/standard_coders_test.py | 93 ++++-----
sdks/python/apache_beam/coders/typecoders.py | 2 +
sdks/python/apache_beam/coders/typecoders_test.py | 3 +
sdks/python/apache_beam/dataframe/convert.py | 14 +-
sdks/python/apache_beam/dataframe/convert_test.py | 4 +
sdks/python/apache_beam/dataframe/doctests.py | 45 ++--
sdks/python/apache_beam/dataframe/doctests_test.py | 1 +
sdks/python/apache_beam/dataframe/expressions.py | 7 +
.../apache_beam/dataframe/expressions_test.py | 1 +
sdks/python/apache_beam/dataframe/frame_base.py | 10 +
.../apache_beam/dataframe/frame_base_test.py | 15 ++
sdks/python/apache_beam/dataframe/frames.py | 191 +++++++----------
sdks/python/apache_beam/dataframe/frames_test.py | 114 +++++-----
sdks/python/apache_beam/dataframe/io.py | 31 ++-
sdks/python/apache_beam/dataframe/io_it_test.py | 1 +
sdks/python/apache_beam/dataframe/io_test.py | 5 +
.../apache_beam/dataframe/pandas_doctests_test.py | 150 ++++++--------
.../dataframe/pandas_top_level_functions.py | 12 +-
sdks/python/apache_beam/dataframe/partitionings.py | 5 +
sdks/python/apache_beam/dataframe/schemas.py | 2 +
sdks/python/apache_beam/dataframe/schemas_test.py | 3 +
sdks/python/apache_beam/dataframe/transforms.py | 24 ++-
.../apache_beam/dataframe/transforms_test.py | 11 +-
sdks/python/apache_beam/examples/avro_nyc_trips.py | 12 +-
.../apache_beam/examples/avro_nyc_trips_test.py | 1 +
.../apache_beam/examples/complete/autocomplete.py | 1 +
.../examples/complete/autocomplete_it_test.py | 1 +
.../apache_beam/examples/complete/distribopt.py | 8 +
.../apache_beam/examples/complete/estimate_pi.py | 2 +
.../examples/complete/estimate_pi_it_test.py | 1 +
.../examples/complete/estimate_pi_test.py | 2 +
.../examples/complete/game/game_stats.py | 5 +
.../examples/complete/game/hourly_team_score.py | 5 +
.../examples/complete/game/leader_board.py | 6 +
.../examples/complete/game/user_score.py | 3 +
.../complete/juliaset/juliaset/juliaset.py | 1 +
.../complete/juliaset/juliaset/juliaset_test.py | 1 +
.../examples/complete/juliaset/setup.py | 1 +
sdks/python/apache_beam/examples/complete/tfidf.py | 1 +
.../apache_beam/examples/complete/tfidf_it_test.py | 1 +
.../apache_beam/examples/complete/tfidf_test.py | 1 +
.../examples/complete/top_wikipedia_sessions.py | 3 +
.../examples/cookbook/bigquery_side_input.py | 1 +
.../examples/cookbook/bigquery_side_input_test.py | 1 +
.../examples/cookbook/bigquery_tornadoes_test.py | 1 +
.../examples/cookbook/bigtableio_it_test.py | 1 +
.../python/apache_beam/examples/cookbook/coders.py | 1 +
.../examples/cookbook/coders_it_test.py | 1 +
.../examples/cookbook/combiners_test.py | 1 +
.../examples/cookbook/custom_ptransform.py | 1 +
.../examples/cookbook/custom_ptransform_it_test.py | 1 +
.../examples/cookbook/datastore_wordcount.py | 2 +
.../examples/cookbook/group_with_coder.py | 2 +
.../examples/cookbook/multiple_output_pardo.py | 2 +
.../examples/dataframe/taxiride_it_test.py | 1 +
.../anomaly_detection_pipeline/main.py | 1 +
.../pipeline/transformations.py | 5 +
.../pipeline/utils.py | 2 +
.../inference/huggingface_language_modeling.py | 1 +
.../inference/huggingface_question_answering.py | 1 +
.../inference/large_language_modeling/main.py | 2 +
.../inference/milk_quality_prediction_windowing.py | 1 +
.../composite_transform.py | 4 +
.../online_clustering/clustering_pipeline/main.py | 1 +
.../pipeline/transformations.py | 4 +
.../pipeline/utils.py | 2 +
.../inference/onnx_sentiment_classification.py | 2 +
...ytorch_image_classification_with_side_inputs.py | 1 +
.../inference/pytorch_image_segmentation.py | 1 +
.../inference/pytorch_language_modeling.py | 27 +--
.../pytorch_model_per_key_image_segmentation.py | 3 +
.../inference/run_inference_side_inputs.py | 5 +
.../pipeline/transformations.py | 4 +
.../sklearn_japanese_housing_regression.py | 1 +
.../inference/sklearn_mnist_classification.py | 1 +
.../inference/tensorflow_imagenet_segmentation.py | 1 +
.../inference/tensorflow_mnist_classification.py | 1 +
.../inference/tensorrt_object_detection.py | 1 +
.../inference/tensorrt_text_classification.py | 2 +
.../inference/tfx_bsl/build_tensorflow_model.py | 1 +
.../tfx_bsl/tensorflow_image_classification.py | 1 +
.../inference/tfx_bsl/tfx_bsl_inference_it_test.py | 1 +
.../inference/vertex_ai_image_classification.py | 1 +
.../examples/inference/vllm_text_completion.py | 1 +
.../inference/xgboost_iris_classification.py | 5 +-
.../examples/ml_transform/ml_transform_it_test.py | 1 +
.../ml_transform/vocab_tfidf_processing.py | 5 +-
.../apache_beam/examples/online_clustering.py | 8 +
.../apache_beam/examples/per_entity_training.py | 4 +
.../apache_beam/examples/snippets/snippets.py | 32 +++
.../snippets_examples_wordcount_debugging.py | 1 +
.../snippets_examples_wordcount_wordcount.py | 1 +
.../apache_beam/examples/snippets/snippets_test.py | 45 +++-
.../aggregation/approximatequantiles_test.py | 5 +-
.../aggregation/approximateunique_test.py | 5 +-
.../transforms/aggregation/batchelements_test.py | 1 +
.../transforms/aggregation/cogroupbykey_test.py | 1 +
.../aggregation/combineglobally_combinefn.py | 1 +
.../transforms/aggregation/combineglobally_test.py | 1 +
.../aggregation/combineperkey_combinefn.py | 1 +
.../combineperkey_side_inputs_singleton.py | 3 +-
.../transforms/aggregation/combineperkey_test.py | 1 +
.../aggregation/combinevalues_combinefn.py | 1 +
.../combinevalues_side_inputs_singleton.py | 3 +-
.../transforms/aggregation/combinevalues_test.py | 1 +
.../snippets/transforms/aggregation/count_test.py | 1 +
.../transforms/aggregation/distinct_test.py | 1 +
.../transforms/aggregation/groupbykey_test.py | 1 +
.../aggregation/groupintobatches_test.py | 1 +
.../snippets/transforms/aggregation/latest_test.py | 1 +
.../snippets/transforms/aggregation/max_test.py | 1 +
.../snippets/transforms/aggregation/mean_test.py | 1 +
.../snippets/transforms/aggregation/min_test.py | 1 +
.../snippets/transforms/aggregation/sample_test.py | 1 +
.../snippets/transforms/aggregation/sum_test.py | 1 +
.../snippets/transforms/aggregation/tolist_test.py | 2 +
.../snippets/transforms/aggregation/top_of.py | 7 +-
.../snippets/transforms/aggregation/top_per_key.py | 7 +-
.../snippets/transforms/aggregation/top_test.py | 1 +
.../transforms/elementwise/enrichment_test.py | 1 +
.../elementwise/filter_side_inputs_dict.py | 3 +-
.../elementwise/filter_side_inputs_iter.py | 3 +-
.../elementwise/filter_side_inputs_singleton.py | 3 +-
.../snippets/transforms/elementwise/filter_test.py | 1 +
.../elementwise/flatmap_side_inputs_singleton.py | 3 +-
.../transforms/elementwise/flatmap_test.py | 1 +
.../snippets/transforms/elementwise/keys_test.py | 1 +
.../snippets/transforms/elementwise/kvswap_test.py | 1 +
.../snippets/transforms/elementwise/map_context.py | 5 +-
.../transforms/elementwise/map_side_inputs_iter.py | 3 +-
.../elementwise/map_side_inputs_singleton.py | 3 +-
.../snippets/transforms/elementwise/map_test.py | 2 +
.../transforms/elementwise/mltransform_test.py | 1 +
.../snippets/transforms/elementwise/pardo_dofn.py | 1 +
.../transforms/elementwise/pardo_dofn_methods.py | 1 +
.../transforms/elementwise/pardo_dofn_params.py | 1 +
.../snippets/transforms/elementwise/pardo_test.py | 2 +
.../transforms/elementwise/partition_test.py | 1 +
.../snippets/transforms/elementwise/regex_test.py | 1 +
.../transforms/elementwise/runinference.py | 1 +
.../transforms/elementwise/runinference_test.py | 2 +
.../transforms/elementwise/tostring_test.py | 1 +
.../snippets/transforms/elementwise/values_test.py | 1 +
.../elementwise/withtimestamps_event_time.py | 24 ++-
.../elementwise/withtimestamps_logical_clock.py | 1 +
.../elementwise/withtimestamps_processing_time.py | 1 +
.../transforms/elementwise/withtimestamps_test.py | 1 +
.../snippets/transforms/other/create_test.py | 1 +
.../snippets/transforms/other/flatten_test.py | 1 +
.../examples/snippets/transforms/other/window.py | 48 +++--
.../snippets/transforms/other/window_test.py | 1 +
sdks/python/apache_beam/examples/snippets/util.py | 1 +
.../apache_beam/examples/snippets/util_test.py | 1 +
sdks/python/apache_beam/examples/sql_taxi.py | 15 +-
.../examples/streaming_wordcount_debugging.py | 2 +
.../streaming_wordcount_debugging_it_test.py | 1 +
.../examples/streaming_wordcount_debugging_test.py | 3 +
.../examples/streaming_wordcount_it_test.py | 1 +
.../apache_beam/examples/windowed_wordcount.py | 1 +
sdks/python/apache_beam/examples/wordcount.py | 1 +
.../apache_beam/examples/wordcount_debugging.py | 3 +
.../apache_beam/examples/wordcount_with_metrics.py | 1 +
.../python/apache_beam/examples/wordcount_xlang.py | 1 +
sdks/python/apache_beam/internal/dill_pickler.py | 5 +-
sdks/python/apache_beam/internal/gcp/auth.py | 1 +
sdks/python/apache_beam/internal/gcp/auth_test.py | 2 +
.../apache_beam/internal/gcp/json_value_test.py | 1 +
.../apache_beam/internal/http_client_test.py | 1 +
sdks/python/apache_beam/internal/metrics/cells.py | 4 +
.../apache_beam/internal/metrics/cells_test.py | 1 +
sdks/python/apache_beam/internal/metrics/metric.py | 8 +-
.../apache_beam/internal/metrics/metric_test.py | 3 +
sdks/python/apache_beam/internal/module_test.py | 7 +
sdks/python/apache_beam/internal/util.py | 8 +-
sdks/python/apache_beam/internal/util_test.py | 1 +
sdks/python/apache_beam/io/avroio.py | 15 +-
sdks/python/apache_beam/io/avroio_test.py | 6 +-
.../apache_beam/io/aws/clients/s3/boto3_client.py | 1 +
.../apache_beam/io/aws/clients/s3/client_test.py | 1 +
.../apache_beam/io/aws/clients/s3/fake_client.py | 2 +
.../apache_beam/io/aws/clients/s3/messages.py | 14 ++
.../python/apache_beam/io/aws/s3filesystem_test.py | 1 +
sdks/python/apache_beam/io/aws/s3io.py | 3 +
sdks/python/apache_beam/io/aws/s3io_test.py | 1 +
.../io/azure/blobstoragefilesystem_test.py | 1 +
sdks/python/apache_beam/io/azure/blobstorageio.py | 5 +
sdks/python/apache_beam/io/components/util.py | 1 +
sdks/python/apache_beam/io/concat_source.py | 2 +
sdks/python/apache_beam/io/concat_source_test.py | 1 +
sdks/python/apache_beam/io/debezium.py | 1 +
.../io/external/generate_sequence_test.py | 1 +
.../io/external/xlang_debeziumio_it_test.py | 5 +-
.../io/external/xlang_jdbcio_it_test.py | 4 +-
.../io/external/xlang_kafkaio_it_test.py | 2 +
.../io/external/xlang_kafkaio_perf_test.py | 4 +
.../io/external/xlang_kinesisio_it_test.py | 2 +
.../io/external/xlang_snowflakeio_it_test.py | 3 +
.../python/apache_beam/io/filebasedio_perf_test.py | 6 +
sdks/python/apache_beam/io/filebasedsink.py | 10 +-
sdks/python/apache_beam/io/filebasedsink_test.py | 4 +
sdks/python/apache_beam/io/filebasedsource.py | 4 +
sdks/python/apache_beam/io/filebasedsource_test.py | 7 +
sdks/python/apache_beam/io/fileio.py | 28 ++-
sdks/python/apache_beam/io/fileio_test.py | 15 +-
sdks/python/apache_beam/io/filesystem.py | 3 +
sdks/python/apache_beam/io/filesystem_test.py | 10 +-
sdks/python/apache_beam/io/filesystemio.py | 5 +
sdks/python/apache_beam/io/filesystemio_test.py | 5 +
sdks/python/apache_beam/io/filesystems_test.py | 2 +
.../flink/flink_streaming_impulse_source_test.py | 1 +
sdks/python/apache_beam/io/gcp/__init__.py | 2 +
.../io/gcp/big_query_query_to_table_it_test.py | 1 +
sdks/python/apache_beam/io/gcp/bigquery.py | 51 +++--
.../apache_beam/io/gcp/bigquery_avro_tools_test.py | 1 +
.../apache_beam/io/gcp/bigquery_file_loads.py | 11 +-
.../apache_beam/io/gcp/bigquery_file_loads_test.py | 20 +-
.../apache_beam/io/gcp/bigquery_io_metadata.py | 1 +
.../io/gcp/bigquery_io_metadata_test.py | 1 +
.../io/gcp/bigquery_io_read_pipeline.py | 1 +
.../apache_beam/io/gcp/bigquery_json_it_test.py | 1 +
.../apache_beam/io/gcp/bigquery_read_internal.py | 9 +
.../apache_beam/io/gcp/bigquery_read_it_test.py | 3 +
.../apache_beam/io/gcp/bigquery_read_perf_test.py | 1 +
.../apache_beam/io/gcp/bigquery_schema_tools.py | 1 +
.../io/gcp/bigquery_schema_tools_test.py | 2 +
sdks/python/apache_beam/io/gcp/bigquery_test.py | 85 ++++----
sdks/python/apache_beam/io/gcp/bigquery_tools.py | 9 +-
.../apache_beam/io/gcp/bigquery_tools_test.py | 19 +-
.../apache_beam/io/gcp/bigquery_write_it_test.py | 6 +-
.../apache_beam/io/gcp/bigquery_write_perf_test.py | 1 +
sdks/python/apache_beam/io/gcp/bigtableio.py | 3 +
.../datastore/v1new/datastore_write_it_pipeline.py | 1 +
.../io/gcp/datastore/v1new/datastoreio.py | 7 +
.../io/gcp/datastore/v1new/datastoreio_test.py | 4 +
.../io/gcp/datastore/v1new/query_splitter.py | 1 +
.../datastore/v1new/rampup_throttling_fn_test.py | 1 +
.../apache_beam/io/gcp/datastore/v1new/types.py | 3 +
.../apache_beam/io/gcp/datastore/v1new/util.py | 2 +
.../io/gcp/datastore/v1new/util_test.py | 1 +
.../apache_beam/io/gcp/experimental/spannerio.py | 45 ++--
.../gcp/experimental/spannerio_read_perf_test.py | 2 +
.../io/gcp/experimental/spannerio_test.py | 2 +
.../gcp/experimental/spannerio_write_perf_test.py | 1 +
.../apache_beam/io/gcp/gcsfilesystem_test.py | 7 +-
sdks/python/apache_beam/io/gcp/gcsio.py | 3 +
sdks/python/apache_beam/io/gcp/gcsio_retry_test.py | 1 +
sdks/python/apache_beam/io/gcp/gcsio_test.py | 3 +
.../apache_beam/io/gcp/healthcare/dicomio.py | 6 +
.../io/gcp/healthcare/dicomio_integration_test.py | 1 +
.../apache_beam/io/gcp/healthcare/dicomio_test.py | 2 +
sdks/python/apache_beam/io/gcp/pubsub.py | 8 +
.../apache_beam/io/gcp/pubsub_io_perf_test.py | 4 +
sdks/python/apache_beam/io/gcp/pubsub_test.py | 14 +-
.../apache_beam/io/gcp/pubsublite/external.py | 2 +
.../apache_beam/io/gcp/pubsublite/proto_api.py | 2 +
sdks/python/apache_beam/io/gcp/spanner.py | 1 +
.../apache_beam/io/gcp/tests/bigquery_matcher.py | 12 +-
.../io/gcp/tests/bigquery_matcher_test.py | 3 +
.../apache_beam/io/gcp/tests/pubsub_matcher.py | 1 +
.../io/gcp/tests/pubsub_matcher_test.py | 1 +
sdks/python/apache_beam/io/gcp/tests/utils_test.py | 13 +-
.../io/gcp/tests/xlang_spannerio_it_test.py | 3 +
sdks/python/apache_beam/io/hadoopfilesystem.py | 3 +
.../python/apache_beam/io/hadoopfilesystem_test.py | 3 +
sdks/python/apache_beam/io/iobase.py | 16 ++
sdks/python/apache_beam/io/iobase_test.py | 4 +
sdks/python/apache_beam/io/jdbc.py | 2 +
sdks/python/apache_beam/io/localfilesystem.py | 2 +
sdks/python/apache_beam/io/localfilesystem_test.py | 7 +-
sdks/python/apache_beam/io/mongodbio.py | 20 +-
sdks/python/apache_beam/io/mongodbio_it_test.py | 1 +
sdks/python/apache_beam/io/mongodbio_test.py | 18 +-
sdks/python/apache_beam/io/parquetio.py | 12 ++
sdks/python/apache_beam/io/parquetio_it_test.py | 2 +
sdks/python/apache_beam/io/parquetio_test.py | 1 +
sdks/python/apache_beam/io/range_trackers.py | 8 +-
sdks/python/apache_beam/io/range_trackers_test.py | 4 +
sdks/python/apache_beam/io/requestresponse.py | 16 ++
.../apache_beam/io/requestresponse_it_test.py | 7 +
sdks/python/apache_beam/io/requestresponse_test.py | 5 +
sdks/python/apache_beam/io/restriction_trackers.py | 3 +
.../apache_beam/io/restriction_trackers_test.py | 2 +
sdks/python/apache_beam/io/source_test_utils.py | 8 +-
.../apache_beam/io/source_test_utils_test.py | 1 +
sdks/python/apache_beam/io/sources_test.py | 1 +
sdks/python/apache_beam/io/textio.py | 4 +
sdks/python/apache_beam/io/textio_test.py | 4 +
sdks/python/apache_beam/io/tfrecordio.py | 6 +
sdks/python/apache_beam/io/tfrecordio_test.py | 6 +
sdks/python/apache_beam/io/utils.py | 1 +
sdks/python/apache_beam/io/utils_test.py | 1 +
sdks/python/apache_beam/io/watermark_estimators.py | 9 +
.../apache_beam/io/watermark_estimators_test.py | 3 +
sdks/python/apache_beam/metrics/cells.py | 16 +-
sdks/python/apache_beam/metrics/cells_test.py | 9 +-
sdks/python/apache_beam/metrics/execution.py | 29 +--
sdks/python/apache_beam/metrics/execution_test.py | 2 +
sdks/python/apache_beam/metrics/metric.py | 7 +
sdks/python/apache_beam/metrics/metric_test.py | 6 +
sdks/python/apache_beam/metrics/metricbase.py | 8 +
.../python/apache_beam/metrics/monitoring_infos.py | 3 +-
.../apache_beam/metrics/monitoring_infos_test.py | 1 +
sdks/python/apache_beam/ml/gcp/cloud_dlp.py | 4 +
.../python/apache_beam/ml/gcp/cloud_dlp_it_test.py | 1 +
sdks/python/apache_beam/ml/gcp/cloud_dlp_test.py | 8 +
.../python/apache_beam/ml/gcp/naturallanguageml.py | 2 +
.../apache_beam/ml/gcp/naturallanguageml_test.py | 1 +
.../ml/gcp/naturallanguageml_test_it.py | 1 +
.../apache_beam/ml/gcp/recommendations_ai.py | 7 +
.../apache_beam/ml/gcp/recommendations_ai_test.py | 3 +
.../apache_beam/ml/gcp/videointelligenceml.py | 4 +
.../apache_beam/ml/gcp/videointelligenceml_test.py | 1 +
sdks/python/apache_beam/ml/gcp/visionml.py | 2 +
sdks/python/apache_beam/ml/gcp/visionml_test.py | 1 +
sdks/python/apache_beam/ml/gcp/visionml_test_it.py | 1 +
sdks/python/apache_beam/ml/inference/base.py | 15 ++
sdks/python/apache_beam/ml/inference/base_test.py | 39 +++-
.../ml/inference/huggingface_inference.py | 22 +-
.../ml/inference/huggingface_inference_it_test.py | 1 +
.../ml/inference/huggingface_inference_test.py | 8 +-
.../ml/inference/onnx_inference_it_test.py | 1 +
.../ml/inference/onnx_inference_test.py | 78 ++++---
.../apache_beam/ml/inference/pytorch_inference.py | 4 +
.../ml/inference/pytorch_inference_it_test.py | 5 +-
.../ml/inference/pytorch_inference_test.py | 94 +++++----
.../apache_beam/ml/inference/sklearn_inference.py | 2 +
.../ml/inference/sklearn_inference_it_test.py | 1 +
.../ml/inference/sklearn_inference_test.py | 8 +-
.../ml/inference/tensorflow_inference.py | 2 +
.../ml/inference/tensorflow_inference_it_test.py | 1 +
.../ml/inference/tensorflow_inference_test.py | 37 ++--
.../apache_beam/ml/inference/tensorrt_inference.py | 2 +
.../ml/inference/tensorrt_inference_test.py | 26 ++-
sdks/python/apache_beam/ml/inference/utils.py | 5 +-
sdks/python/apache_beam/ml/inference/utils_test.py | 1 +
.../ml/inference/vertex_ai_inference.py | 1 +
.../ml/inference/vertex_ai_inference_it_test.py | 1 +
.../ml/inference/vertex_ai_inference_test.py | 2 +
.../apache_beam/ml/inference/vllm_inference.py | 3 +
.../apache_beam/ml/inference/xgboost_inference.py | 5 +
.../ml/inference/xgboost_inference_it_test.py | 1 +
.../ml/inference/xgboost_inference_test.py | 1 +
sdks/python/apache_beam/ml/rag/chunking/base.py | 1 +
.../apache_beam/ml/rag/chunking/base_test.py | 5 +
.../apache_beam/ml/rag/chunking/langchain.py | 2 +
.../apache_beam/ml/rag/chunking/langchain_test.py | 4 +-
.../apache_beam/ml/rag/embeddings/base_test.py | 1 +
.../apache_beam/ml/rag/embeddings/huggingface.py | 1 +
.../ml/rag/embeddings/huggingface_test.py | 1 +
sdks/python/apache_beam/ml/transforms/base.py | 17 +-
sdks/python/apache_beam/ml/transforms/base_test.py | 31 ++-
.../ml/transforms/embeddings/huggingface.py | 8 +-
.../ml/transforms/embeddings/huggingface_test.py | 3 +
.../ml/transforms/embeddings/tensorflow_hub.py | 3 +
.../transforms/embeddings/tensorflow_hub_test.py | 3 +
.../ml/transforms/embeddings/vertex_ai.py | 4 +
.../ml/transforms/embeddings/vertex_ai_test.py | 2 +
sdks/python/apache_beam/ml/transforms/handlers.py | 4 +
.../apache_beam/ml/transforms/handlers_test.py | 4 +
sdks/python/apache_beam/ml/transforms/tft.py | 15 ++
sdks/python/apache_beam/ml/transforms/tft_test.py | 31 ++-
sdks/python/apache_beam/ml/transforms/utils.py | 1 +
.../python/apache_beam/options/pipeline_options.py | 31 ++-
.../apache_beam/options/pipeline_options_test.py | 14 ++
.../options/pipeline_options_validator.py | 4 +-
.../options/pipeline_options_validator_test.py | 9 +
sdks/python/apache_beam/options/value_provider.py | 4 +
.../apache_beam/options/value_provider_test.py | 18 ++
sdks/python/apache_beam/pipeline.py | 39 ++--
sdks/python/apache_beam/pipeline_test.py | 77 ++++++-
sdks/python/apache_beam/pvalue.py | 25 ++-
sdks/python/apache_beam/pvalue_test.py | 3 +
sdks/python/apache_beam/runners/common.py | 229 +++++++++++----------
sdks/python/apache_beam/runners/common_test.py | 31 ++-
.../python/apache_beam/runners/dask/dask_runner.py | 2 +
.../apache_beam/runners/dask/dask_runner_test.py | 23 ++-
sdks/python/apache_beam/runners/dask/overrides.py | 11 +
.../runners/dask/transform_evaluator.py | 6 +
.../dataflow/dataflow_exercise_metrics_pipeline.py | 11 +-
.../dataflow_exercise_metrics_pipeline_test.py | 1 +
...dataflow_exercise_streaming_metrics_pipeline.py | 1 +
...low_exercise_streaming_metrics_pipeline_test.py | 1 +
.../runners/dataflow/dataflow_job_service.py | 4 +-
.../runners/dataflow/dataflow_job_service_test.py | 2 +
.../runners/dataflow/dataflow_metrics.py | 5 +-
.../runners/dataflow/dataflow_metrics_test.py | 10 +-
.../runners/dataflow/dataflow_runner.py | 12 +-
.../runners/dataflow/dataflow_runner_test.py | 13 +-
.../runners/dataflow/internal/apiclient.py | 22 +-
.../runners/dataflow/internal/apiclient_test.py | 3 +
.../clients/cloudbuild/cloudbuild_v1_client.py | 168 +++++----------
.../clients/cloudbuild/cloudbuild_v1_messages.py | 38 ++++
.../internal/clients/dataflow/message_matchers.py | 2 +
.../clients/dataflow/message_matchers_test.py | 1 +
.../runners/dataflow/ptransform_overrides.py | 1 +
.../runners/dataflow/template_runner_test.py | 1 +
.../runners/dataflow/test_dataflow_runner.py | 1 +
.../apache_beam/runners/direct/bundle_factory.py | 3 +
sdks/python/apache_beam/runners/direct/clock.py | 3 +
.../direct/consumer_tracking_pipeline_visitor.py | 1 +
.../consumer_tracking_pipeline_visitor_test.py | 10 +-
.../apache_beam/runners/direct/direct_metrics.py | 25 ++-
.../apache_beam/runners/direct/direct_runner.py | 18 ++
.../runners/direct/direct_runner_test.py | 7 +
.../apache_beam/runners/direct/direct_userstate.py | 6 +
.../runners/direct/evaluation_context.py | 10 +-
sdks/python/apache_beam/runners/direct/executor.py | 11 +
.../runners/direct/helper_transforms.py | 3 +
.../runners/direct/sdf_direct_runner.py | 17 ++
.../runners/direct/sdf_direct_runner_test.py | 5 +
.../runners/direct/test_direct_runner.py | 1 +
.../apache_beam/runners/direct/test_stream_impl.py | 2 +
.../runners/direct/transform_evaluator.py | 15 +-
sdks/python/apache_beam/runners/direct/util.py | 3 +
.../runners/direct/watermark_manager.py | 1 +
.../runners/interactive/augmented_pipeline.py | 1 +
.../runners/interactive/augmented_pipeline_test.py | 2 +
.../runners/interactive/background_caching_job.py | 1 +
.../interactive/background_caching_job_test.py | 28 +--
.../runners/interactive/cache_manager.py | 15 +-
.../interactive/caching/expression_cache.py | 1 +
.../interactive/caching/expression_cache_test.py | 1 +
.../runners/interactive/caching/read_cache.py | 6 +-
.../runners/interactive/caching/read_cache_test.py | 1 +
.../runners/interactive/caching/reify.py | 2 +
.../runners/interactive/caching/streaming_cache.py | 6 +
.../interactive/caching/streaming_cache_test.py | 1 +
.../runners/interactive/caching/write_cache.py | 20 +-
.../interactive/caching/write_cache_test.py | 1 +
.../dataproc/dataproc_cluster_manager.py | 1 +
.../dataproc/dataproc_cluster_manager_test.py | 6 +
.../runners/interactive/display/display_manager.py | 1 +
.../display/interactive_pipeline_graph.py | 1 +
.../interactive/display/pcoll_visualization.py | 2 +
.../display/pcoll_visualization_test.py | 1 +
.../runners/interactive/display/pipeline_graph.py | 2 +
.../interactive/display/pipeline_graph_renderer.py | 4 +
.../interactive/display/pipeline_graph_test.py | 1 +
.../runners/interactive/interactive_beam.py | 2 +
.../runners/interactive/interactive_beam_test.py | 5 +
.../runners/interactive/interactive_environment.py | 4 +-
.../interactive/interactive_environment_test.py | 4 +
.../runners/interactive/interactive_runner.py | 3 +
.../runners/interactive/interactive_runner_test.py | 10 +
.../messaging/interactive_environment_inspector.py | 1 +
.../interactive_environment_inspector_test.py | 1 +
.../interactive/non_interactive_runner_test.py | 5 +
.../runners/interactive/options/capture_control.py | 1 +
.../interactive/options/capture_control_test.py | 4 +
.../interactive/options/capture_limiters.py | 6 +
.../interactive/options/capture_limiters_test.py | 1 +
.../interactive/options/interactive_options.py | 1 +
.../runners/interactive/pipeline_fragment.py | 5 +
.../runners/interactive/pipeline_fragment_test.py | 1 +
.../runners/interactive/pipeline_instrument.py | 23 ++-
.../interactive/pipeline_instrument_test.py | 8 +
.../runners/interactive/recording_manager.py | 14 +-
.../runners/interactive/recording_manager_test.py | 7 +
.../runners/interactive/sql/beam_sql_magics.py | 3 +
.../interactive/sql/beam_sql_magics_test.py | 25 +--
.../runners/interactive/sql/sql_chain.py | 5 +-
.../runners/interactive/sql/sql_chain_test.py | 1 +
.../apache_beam/runners/interactive/sql/utils.py | 6 +-
.../runners/interactive/sql/utils_test.py | 2 +
.../testing/integration/notebook_executor.py | 2 +
.../interactive/testing/integration/screen_diff.py | 1 +
.../testing/integration/tests/screen_diff_test.py | 2 +
.../runners/interactive/testing/mock_env.py | 2 +
.../runners/interactive/testing/mock_ipython.py | 2 +
.../interactive/testing/test_cache_manager.py | 3 +
.../runners/interactive/user_pipeline_tracker.py | 1 +
.../interactive/user_pipeline_tracker_test.py | 1 +
.../apache_beam/runners/interactive/utils.py | 4 +
.../apache_beam/runners/interactive/utils_test.py | 9 +
sdks/python/apache_beam/runners/job/manager.py | 1 +
.../python/apache_beam/runners/pipeline_context.py | 3 +
.../apache_beam/runners/pipeline_context_test.py | 1 +
.../runners/portability/abstract_job_service.py | 15 +-
.../runners/portability/artifact_service.py | 6 +-
.../runners/portability/artifact_service_test.py | 3 +
.../runners/portability/expansion_service.py | 16 +-
.../runners/portability/expansion_service_test.py | 19 ++
.../runners/portability/flink_runner.py | 1 +
.../runners/portability/flink_runner_test.py | 5 +-
.../portability/flink_uber_jar_job_server.py | 3 +
.../portability/flink_uber_jar_job_server_test.py | 1 +
.../runners/portability/fn_api_runner/execution.py | 17 +-
.../runners/portability/fn_api_runner/fn_runner.py | 36 ++--
.../portability/fn_api_runner/fn_runner_test.py | 131 ++++++++++--
.../portability/fn_api_runner/translations.py | 46 +++--
.../portability/fn_api_runner/translations_test.py | 29 ++-
.../portability/fn_api_runner/trigger_manager.py | 5 +
.../fn_api_runner/trigger_manager_test.py | 16 +-
.../portability/fn_api_runner/watermark_manager.py | 3 +
.../portability/fn_api_runner/worker_handlers.py | 165 ++++++++-------
.../fn_api_runner/worker_handlers_test.py | 1 +
.../apache_beam/runners/portability/job_server.py | 6 +
.../runners/portability/job_server_test.py | 2 +
.../runners/portability/local_job_service.py | 19 +-
.../runners/portability/local_job_service_test.py | 2 +
.../runners/portability/portable_runner.py | 9 +-
.../runners/portability/portable_runner_test.py | 8 +
.../runners/portability/prism_runner.py | 5 +-
.../runners/portability/sdk_container_builder.py | 3 +
.../portability/sdk_container_builder_test.py | 4 +
.../portability/spark_java_job_server_test.py | 2 +
.../runners/portability/spark_runner.py | 1 +
.../portability/spark_uber_jar_job_server.py | 2 +
.../portability/spark_uber_jar_job_server_test.py | 1 +
.../apache_beam/runners/portability/stager.py | 4 +-
.../apache_beam/runners/portability/stager_test.py | 2 +
sdks/python/apache_beam/runners/render.py | 11 +-
sdks/python/apache_beam/runners/render_test.py | 2 +
sdks/python/apache_beam/runners/runner.py | 2 +
sdks/python/apache_beam/runners/sdf_utils.py | 5 +
sdks/python/apache_beam/runners/sdf_utils_test.py | 3 +
sdks/python/apache_beam/runners/trivial_runner.py | 10 +-
.../apache_beam/runners/trivial_runner_test.py | 1 +
.../apache_beam/runners/worker/bundle_processor.py | 58 ++++--
.../runners/worker/bundle_processor_test.py | 14 ++
.../apache_beam/runners/worker/data_plane.py | 8 +-
.../apache_beam/runners/worker/data_plane_test.py | 1 +
.../apache_beam/runners/worker/data_sampler.py | 11 +-
.../runners/worker/data_sampler_test.py | 2 +
.../apache_beam/runners/worker/log_handler.py | 4 +-
.../apache_beam/runners/worker/log_handler_test.py | 9 +-
sdks/python/apache_beam/runners/worker/logger.py | 2 +
.../apache_beam/runners/worker/logger_test.py | 3 +
.../apache_beam/runners/worker/opcounters.py | 19 +-
.../apache_beam/runners/worker/opcounters_test.py | 4 +
.../apache_beam/runners/worker/operation_specs.py | 23 +--
.../apache_beam/runners/worker/operations.py | 123 ++++++-----
.../apache_beam/runners/worker/sdk_worker.py | 27 ++-
.../apache_beam/runners/worker/sdk_worker_main.py | 3 +-
.../runners/worker/sdk_worker_main_test.py | 10 +-
.../apache_beam/runners/worker/sdk_worker_test.py | 62 +++---
.../apache_beam/runners/worker/sideinputs.py | 4 +
.../apache_beam/runners/worker/sideinputs_test.py | 7 +
.../apache_beam/runners/worker/statecache.py | 4 +
.../apache_beam/runners/worker/statecache_test.py | 39 ++--
.../apache_beam/runners/worker/statesampler.py | 1 +
.../runners/worker/statesampler_slow.py | 2 +
.../runners/worker/worker_id_interceptor_test.py | 1 +
.../apache_beam/runners/worker/worker_pool_main.py | 2 +
.../apache_beam/runners/worker/worker_status.py | 5 +-
.../runners/worker/worker_status_test.py | 4 +
.../testing/analyzers/load_test_perf_analysis.py | 1 +
.../apache_beam/testing/analyzers/perf_analysis.py | 4 +-
.../testing/analyzers/perf_analysis_test.py | 2 +
.../testing/analyzers/perf_analysis_utils.py | 2 +
.../testing/benchmarks/chicago_taxi/preprocess.py | 1 +
.../benchmarks/cloudml/cloudml_benchmark_test.py | 1 +
.../benchmarks/cloudml/criteo_tft/criteo.py | 1 +
.../benchmarks/cloudml/pipelines/workflow.py | 6 +-
.../pytorch_image_classification_benchmarks.py | 1 +
.../pytorch_language_modeling_benchmarks.py | 1 +
...nsorflow_mnist_classification_cost_benchmark.py | 1 +
.../benchmarks/nexmark/models/auction_bid.py | 1 +
.../benchmarks/nexmark/models/nexmark_model.py | 3 +
.../testing/benchmarks/nexmark/monitor.py | 2 +
.../testing/benchmarks/nexmark/nexmark_perf.py | 1 +
.../testing/benchmarks/nexmark/nexmark_util.py | 5 +
.../nexmark/queries/nexmark_query_util.py | 7 +
.../testing/benchmarks/nexmark/queries/query0.py | 1 +
.../testing/benchmarks/nexmark/queries/query1.py | 5 +-
.../testing/benchmarks/nexmark/queries/query10.py | 4 +
.../testing/benchmarks/nexmark/queries/query3.py | 6 +-
.../testing/benchmarks/nexmark/queries/query4.py | 1 +
.../testing/benchmarks/nexmark/queries/query5.py | 1 +
.../testing/benchmarks/nexmark/queries/query6.py | 1 +
.../testing/benchmarks/nexmark/queries/query7.py | 1 +
.../testing/benchmarks/nexmark/queries/query8.py | 1 +
.../benchmarks/nexmark/queries/winning_bids.py | 5 +
.../testing/benchmarks/wordcount/wordcount.py | 1 +
.../apache_beam/testing/datatype_inference.py | 13 +-
.../apache_beam/testing/datatype_inference_test.py | 6 +-
.../python/apache_beam/testing/extra_assertions.py | 1 +
.../apache_beam/testing/extra_assertions_test.py | 1 +
.../testing/load_tests/co_group_by_key_test.py | 1 +
.../apache_beam/testing/load_tests/combine_test.py | 2 +
.../testing/load_tests/dataflow_cost_benchmark.py | 1 +
.../testing/load_tests/group_by_key_test.py | 1 +
.../apache_beam/testing/load_tests/load_test.py | 2 +
.../testing/load_tests/load_test_metrics_utils.py | 21 +-
.../testing/load_tests/microbenchmarks_test.py | 1 +
.../apache_beam/testing/load_tests/pardo_test.py | 4 +
.../testing/load_tests/sideinput_test.py | 5 +
.../apache_beam/testing/metric_result_matchers.py | 2 +
.../testing/metric_result_matchers_test.py | 1 +
.../apache_beam/testing/pipeline_verifiers.py | 2 +
.../apache_beam/testing/pipeline_verifiers_test.py | 1 +
.../apache_beam/testing/synthetic_pipeline.py | 19 +-
.../apache_beam/testing/synthetic_pipeline_test.py | 1 +
.../apache_beam/testing/test_pipeline_test.py | 2 +
sdks/python/apache_beam/testing/test_stream.py | 8 +
.../apache_beam/testing/test_stream_it_test.py | 4 +
.../apache_beam/testing/test_stream_service.py | 1 +
.../testing/test_stream_service_test.py | 2 +
.../python/apache_beam/testing/test_stream_test.py | 12 ++
sdks/python/apache_beam/testing/test_utils.py | 3 +
sdks/python/apache_beam/testing/test_utils_test.py | 1 +
sdks/python/apache_beam/testing/util.py | 10 +
sdks/python/apache_beam/testing/util_test.py | 1 +
.../apache_beam/tools/coders_microbenchmark.py | 4 +
.../tools/fn_api_runner_microbenchmark.py | 1 +
.../apache_beam/tools/microbenchmarks_test.py | 1 +
.../tools/runtime_type_check_microbenchmark.py | 4 +
.../apache_beam/tools/teststream_microbenchmark.py | 2 +
sdks/python/apache_beam/tools/utils.py | 1 +
.../apache_beam/transforms/batch_dofn_test.py | 18 +-
.../transforms/combinefn_lifecycle_test.py | 2 +
sdks/python/apache_beam/transforms/combiners.py | 57 ++++-
.../apache_beam/transforms/combiners_test.py | 21 ++
sdks/python/apache_beam/transforms/core.py | 105 +++++++---
sdks/python/apache_beam/transforms/core_test.py | 20 ++
.../python/apache_beam/transforms/create_source.py | 1 +
sdks/python/apache_beam/transforms/create_test.py | 3 +
sdks/python/apache_beam/transforms/cy_combiners.py | 12 ++
.../dataflow_distribution_counter_test.py | 1 +
sdks/python/apache_beam/transforms/deduplicate.py | 3 +
.../apache_beam/transforms/deduplicate_test.py | 1 +
sdks/python/apache_beam/transforms/display.py | 6 +-
sdks/python/apache_beam/transforms/display_test.py | 15 ++
.../apache_beam/transforms/dofn_lifecycle_test.py | 3 +
sdks/python/apache_beam/transforms/enrichment.py | 2 +
.../transforms/enrichment_handlers/bigquery.py | 1 +
.../enrichment_handlers/bigquery_test.py | 1 +
.../enrichment_handlers/bigtable_it_test.py | 2 +
.../enrichment_handlers/bigtable_test.py | 1 +
.../enrichment_handlers/feast_feature_store.py | 1 +
.../feast_feature_store_it_test.py | 1 +
.../feast_feature_store_test.py | 1 +
.../enrichment_handlers/vertex_ai_feature_store.py | 2 +
.../vertex_ai_feature_store_it_test.py | 2 +
.../vertex_ai_feature_store_test.py | 1 +
.../apache_beam/transforms/enrichment_it_test.py | 2 +
.../apache_beam/transforms/enrichment_test.py | 1 +
sdks/python/apache_beam/transforms/environments.py | 91 +++++---
.../apache_beam/transforms/environments_test.py | 2 +
.../apache_beam/transforms/error_handling.py | 3 +
.../apache_beam/transforms/error_handling_test.py | 3 +
sdks/python/apache_beam/transforms/external.py | 41 ++--
.../apache_beam/transforms/external_it_test.py | 3 +
.../python/apache_beam/transforms/external_java.py | 1 +
.../python/apache_beam/transforms/external_test.py | 15 ++
.../transforms/external_transform_provider.py | 1 +
.../external_transform_provider_it_test.py | 2 +
.../fully_qualified_named_transform_test.py | 15 +-
.../apache_beam/transforms/periodicsequence.py | 4 +
.../transforms/periodicsequence_it_test.py | 3 +
.../transforms/periodicsequence_test.py | 1 +
sdks/python/apache_beam/transforms/ptransform.py | 48 +++--
.../apache_beam/transforms/ptransform_test.py | 119 ++++++++---
.../apache_beam/transforms/resources_test.py | 1 +
sdks/python/apache_beam/transforms/sideinputs.py | 2 +
.../apache_beam/transforms/sideinputs_test.py | 39 ++--
sdks/python/apache_beam/transforms/sql_test.py | 4 +-
sdks/python/apache_beam/transforms/stats.py | 9 +
sdks/python/apache_beam/transforms/stats_test.py | 1 +
.../transforms/timestamped_value_type_test.py | 1 +
sdks/python/apache_beam/transforms/timeutil.py | 6 +
.../transforms_keyword_only_args_test.py | 2 +
sdks/python/apache_beam/transforms/trigger.py | 29 ++-
sdks/python/apache_beam/transforms/trigger_test.py | 109 ++++++----
sdks/python/apache_beam/transforms/userstate.py | 13 ++
.../apache_beam/transforms/userstate_test.py | 25 +++
sdks/python/apache_beam/transforms/util.py | 31 ++-
sdks/python/apache_beam/transforms/util_test.py | 83 +++++---
.../transforms/validate_runner_xlang_test.py | 1 +
sdks/python/apache_beam/transforms/window.py | 16 +-
sdks/python/apache_beam/transforms/window_test.py | 5 +
.../transforms/write_ptransform_test.py | 1 +
.../typehints/arrow_type_compatibility.py | 21 +-
.../typehints/arrow_type_compatibility_test.py | 42 ++--
sdks/python/apache_beam/typehints/batch.py | 8 +-
sdks/python/apache_beam/typehints/batch_test.py | 2 +
sdks/python/apache_beam/typehints/decorators.py | 15 +-
.../apache_beam/typehints/decorators_test.py | 26 +++
.../typehints/intrinsic_one_ops_test.py | 1 +
.../typehints/native_type_compatibility.py | 5 +-
.../typehints/native_type_compatibility_test.py | 1 +
sdks/python/apache_beam/typehints/opcodes.py | 1 +
.../typehints/pandas_type_compatibility.py | 13 +-
.../typehints/pandas_type_compatibility_test.py | 28 +--
.../typehints/pytorch_type_compatibility.py | 6 +-
.../typehints/pytorch_type_compatibility_test.py | 2 +
sdks/python/apache_beam/typehints/row_type.py | 5 +-
.../apache_beam/typehints/schema_registry.py | 1 +
sdks/python/apache_beam/typehints/schemas.py | 15 ++
sdks/python/apache_beam/typehints/schemas_test.py | 23 +--
.../apache_beam/typehints/sharded_key_type.py | 1 +
.../apache_beam/typehints/sharded_key_type_test.py | 1 +
.../apache_beam/typehints/testing/strategies.py | 2 +
.../apache_beam/typehints/trivial_inference.py | 8 +-
.../typehints/trivial_inference_test.py | 104 +++++-----
sdks/python/apache_beam/typehints/typecheck.py | 5 +
.../python/apache_beam/typehints/typecheck_test.py | 11 +
.../apache_beam/typehints/typed_pipeline_test.py | 86 ++++++++
sdks/python/apache_beam/typehints/typehints.py | 34 ++-
.../python/apache_beam/typehints/typehints_test.py | 86 ++++++++
sdks/python/apache_beam/utils/annotations.py | 1 +
sdks/python/apache_beam/utils/counters.py | 2 +
sdks/python/apache_beam/utils/counters_test.py | 3 +
sdks/python/apache_beam/utils/histogram.py | 4 +
sdks/python/apache_beam/utils/histogram_test.py | 1 +
.../apache_beam/utils/interactive_utils_test.py | 1 +
.../apache_beam/utils/multi_process_shared.py | 4 +
.../apache_beam/utils/multi_process_shared_test.py | 4 +
sdks/python/apache_beam/utils/plugin.py | 2 +
sdks/python/apache_beam/utils/processes_test.py | 4 +
sdks/python/apache_beam/utils/profiler_test.py | 1 +
sdks/python/apache_beam/utils/proto_utils_test.py | 1 +
sdks/python/apache_beam/utils/python_callable.py | 1 +
.../apache_beam/utils/python_callable_test.py | 1 +
sdks/python/apache_beam/utils/retry.py | 4 +
sdks/python/apache_beam/utils/retry_test.py | 5 +
sdks/python/apache_beam/utils/sharded_key.py | 1 +
sdks/python/apache_beam/utils/shared.py | 2 +
sdks/python/apache_beam/utils/shared_test.py | 5 +
sdks/python/apache_beam/utils/subprocess_server.py | 2 +
.../apache_beam/utils/subprocess_server_test.py | 3 +
.../apache_beam/utils/thread_pool_executor.py | 4 +
.../apache_beam/utils/thread_pool_executor_test.py | 1 +
sdks/python/apache_beam/utils/timestamp.py | 2 +
sdks/python/apache_beam/utils/timestamp_test.py | 2 +
sdks/python/apache_beam/utils/urns.py | 7 +-
sdks/python/apache_beam/utils/windowed_value.py | 5 +
.../apache_beam/utils/windowed_value_test.py | 3 +
.../yaml/examples/testing/examples_test.py | 5 +-
sdks/python/apache_beam/yaml/generate_yaml_docs.py | 3 +-
sdks/python/apache_beam/yaml/json_utils.py | 8 +-
sdks/python/apache_beam/yaml/main_test.py | 1 +
sdks/python/apache_beam/yaml/options.py | 1 +
sdks/python/apache_beam/yaml/readme_test.py | 10 +-
sdks/python/apache_beam/yaml/yaml_combine.py | 1 +
sdks/python/apache_beam/yaml/yaml_combine_test.py | 1 +
.../apache_beam/yaml/yaml_enrichment_test.py | 2 +
sdks/python/apache_beam/yaml/yaml_errors.py | 2 +
sdks/python/apache_beam/yaml/yaml_io.py | 6 +-
sdks/python/apache_beam/yaml/yaml_io_test.py | 10 +-
sdks/python/apache_beam/yaml/yaml_join_test.py | 6 +-
sdks/python/apache_beam/yaml/yaml_mapping.py | 9 +-
sdks/python/apache_beam/yaml/yaml_mapping_test.py | 1 +
sdks/python/apache_beam/yaml/yaml_ml.py | 3 +
sdks/python/apache_beam/yaml/yaml_ml_test.py | 1 +
sdks/python/apache_beam/yaml/yaml_provider.py | 38 ++--
.../apache_beam/yaml/yaml_provider_unit_test.py | 11 +-
sdks/python/apache_beam/yaml/yaml_transform.py | 13 +-
.../apache_beam/yaml/yaml_transform_scope_test.py | 4 +
.../python/apache_beam/yaml/yaml_transform_test.py | 21 +-
.../apache_beam/yaml/yaml_transform_unit_test.py | 2 +
sdks/python/apache_beam/yaml/yaml_udf_test.py | 9 +-
sdks/python/apache_beam/yaml/yaml_utils.py | 6 +-
sdks/python/apache_beam/yaml/yaml_utils_test.py | 1 +
sdks/python/tox.ini | 4 +-
768 files changed, 5092 insertions(+), 2017 deletions(-)
diff --git a/sdks/python/apache_beam/coders/avro_record.py b/sdks/python/apache_beam/coders/avro_record.py
index c9ed26d34eb..c16c922591e 100644
--- a/sdks/python/apache_beam/coders/avro_record.py
+++ b/sdks/python/apache_beam/coders/avro_record.py
@@ -24,6 +24,7 @@ __all__ = ['AvroRecord']
class AvroRecord(object):
"""Simple wrapper class for dictionary records."""
+
def __init__(self, value):
self.record = value
diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py
index 5dff3505290..9292fd5e25d 100644
--- a/sdks/python/apache_beam/coders/coder_impl.py
+++ b/sdks/python/apache_beam/coders/coder_impl.py
@@ -125,6 +125,7 @@ Observables = List[Tuple[observable.ObservableMixin, 'CoderImpl']]
class CoderImpl(object):
"""For internal use only; no backwards-compatibility guarantees."""
+
def encode_to_stream(self, value, stream, nested):
# type: (Any, create_OutputStream, bool) -> None
@@ -211,6 +212,7 @@ class SimpleCoderImpl(CoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
Subclass of CoderImpl implementing stream methods using encode/decode."""
+
def encode_to_stream(self, value, stream, nested):
# type: (Any, create_OutputStream, bool) -> None
@@ -228,6 +230,7 @@ class StreamCoderImpl(CoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
Subclass of CoderImpl implementing encode/decode using stream methods."""
+
def encode(self, value):
# type: (Any) -> bytes
out = create_OutputStream()
@@ -255,6 +258,7 @@ class CallbackCoderImpl(CoderImpl):
This is the default implementation used if Coder._get_impl()
is not overwritten.
"""
+
def __init__(self, encoder, decoder, size_estimator=None):
self._encoder = encoder
self._decoder = decoder
@@ -297,6 +301,7 @@ class CallbackCoderImpl(CoderImpl):
class ProtoCoderImpl(SimpleCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def __init__(self, proto_message_type):
self.proto_message_type = proto_message_type
@@ -311,12 +316,14 @@ class ProtoCoderImpl(SimpleCoderImpl):
class DeterministicProtoCoderImpl(ProtoCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def encode(self, value):
return value.SerializePartialToString(deterministic=True)
class ProtoPlusCoderImpl(SimpleCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def __init__(self, proto_plus_type):
# type: (Type[proto.Message]) -> None
self.proto_plus_type = proto_plus_type
@@ -356,6 +363,7 @@ _ITERABLE_LIKE_TYPES = set() # type: Set[Type]
class FastPrimitivesCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def __init__(
self, fallback_coder_impl, requires_deterministic_step_label=None):
self.fallback_coder_impl = fallback_coder_impl
@@ -610,6 +618,7 @@ class BytesCoderImpl(CoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
A coder for bytes/str objects."""
+
def encode_to_stream(self, value, out, nested):
# type: (bytes, create_OutputStream, bool) -> None
@@ -636,6 +645,7 @@ class BooleanCoderImpl(CoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
A coder for bool objects."""
+
def encode_to_stream(self, value, out, nested):
out.write_byte(1 if value else 0)
@@ -675,12 +685,12 @@ class MapCoderImpl(StreamCoderImpl):
attribute values.
A coder for typing.Mapping objects."""
+
def __init__(
self,
key_coder, # type: CoderImpl
value_coder, # type: CoderImpl
- is_deterministic = False
- ):
+ is_deterministic=False):
self._key_coder = key_coder
self._value_coder = value_coder
self._is_deterministic = is_deterministic
@@ -760,6 +770,7 @@ class NullableCoderImpl(StreamCoderImpl):
class BigEndianShortCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def encode_to_stream(self, value, out, nested):
# type: (int, create_OutputStream, bool) -> None
out.write_bigendian_int16(value)
@@ -776,6 +787,7 @@ class BigEndianShortCoderImpl(StreamCoderImpl):
class SinglePrecisionFloatCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def encode_to_stream(self, value, out, nested):
# type: (float, create_OutputStream, bool) -> None
out.write_bigendian_float(value)
@@ -792,6 +804,7 @@ class SinglePrecisionFloatCoderImpl(StreamCoderImpl):
class FloatCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def encode_to_stream(self, value, out, nested):
# type: (float, create_OutputStream, bool) -> None
out.write_bigendian_double(value)
@@ -863,6 +876,7 @@ class TimestampCoderImpl(StreamCoderImpl):
that of the Java SDK InstantCoder.
https://github.com/apache/beam/blob/f5029b4f0dfff404310b2ef55e2632bbacc7b04f/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/InstantCoder.java#L79
"""
+
def encode_to_stream(self, value, out, nested):
# type: (Timestamp, create_OutputStream, bool) -> None
millis = value.micros // 1000
@@ -889,6 +903,7 @@ class TimestampCoderImpl(StreamCoderImpl):
class TimerCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def __init__(self, key_coder_impl, window_coder_impl):
self._timestamp_coder_impl = TimestampCoderImpl()
self._boolean_coder_impl = BooleanCoderImpl()
@@ -947,6 +962,7 @@ class VarIntCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
A coder for int objects."""
+
def encode_to_stream(self, value, out, nested):
# type: (int, create_OutputStream, bool) -> None
out.write_var_int64(value)
@@ -978,6 +994,7 @@ class SingletonCoderImpl(CoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
A coder that always encodes exactly one value."""
+
def __init__(self, value):
self._value = value
@@ -1005,6 +1022,7 @@ class AbstractComponentCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
CoderImpl for coders that are comprised of several component coders."""
+
def __init__(self, coder_impls):
for c in coder_impls:
assert isinstance(c, CoderImpl), c
@@ -1030,8 +1048,8 @@ class AbstractComponentCoderImpl(StreamCoderImpl):
# type: (create_InputStream, bool) -> Any
return self._construct_from_components([
c.decode_from_stream(
- in_stream, nested or i + 1 < len(self._coder_impls)) for i,
- c in enumerate(self._coder_impls)
+ in_stream, nested or i + 1 < len(self._coder_impls))
+ for i, c in enumerate(self._coder_impls)
])
def estimate_size(self, value, nested=False):
@@ -1061,6 +1079,7 @@ class AbstractComponentCoderImpl(StreamCoderImpl):
class AvroCoderImpl(SimpleCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def __init__(self, schema):
self.parsed_schema = parse_schema(json.loads(schema))
@@ -1077,6 +1096,7 @@ class AvroCoderImpl(SimpleCoderImpl):
class TupleCoderImpl(AbstractComponentCoderImpl):
"""A coder for tuple objects."""
+
def _extract_components(self, value):
return tuple(value)
@@ -1085,6 +1105,7 @@ class TupleCoderImpl(AbstractComponentCoderImpl):
class _ConcatSequence(object):
+
def __init__(self, head, tail):
# type: (Iterable[Any], Iterable[Any]) -> None
self._head = head
@@ -1277,12 +1298,14 @@ class TupleSequenceCoderImpl(SequenceCoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
A coder for homogeneous tuple objects."""
+
def _construct_from_sequence(self, components):
return tuple(components)
class _AbstractIterable(object):
"""Wraps an iterable hiding methods that might not always be available."""
+
def __init__(self, contents):
self._contents = contents
@@ -1315,6 +1338,7 @@ class IterableCoderImpl(SequenceCoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
A coder for homogeneous iterable objects."""
+
def __init__(self, *args, use_abstract_iterable=None, **kwargs):
super().__init__(*args, **kwargs)
if use_abstract_iterable is None:
@@ -1332,6 +1356,7 @@ class ListCoderImpl(SequenceCoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
A coder for homogeneous list objects."""
+
def _construct_from_sequence(self, components):
return components if isinstance(components, list) else list(components)
@@ -1360,6 +1385,7 @@ class PaneInfoCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
Coder for a PaneInfo descriptor."""
+
def _choose_encoding(self, value):
if ((value._index == 0 and value._nonspeculative_index == 0) or
value._timing == PaneInfoTiming_UNKNOWN):
@@ -1422,6 +1448,7 @@ class PaneInfoCoderImpl(StreamCoderImpl):
class _OrderedUnionCoderImpl(StreamCoderImpl):
+
def __init__(self, coder_impl_types, fallback_coder_impl):
assert len(coder_impl_types) < 128
self._types, self._coder_impls = zip(*coder_impl_types)
@@ -1555,6 +1582,7 @@ class ParamWindowedValueCoderImpl(WindowedValueCoderImpl):
encoding, and uses the supplied parameterized timestamp, windows
and pane info values during decoding when reconstructing the windowed
value."""
+
def __init__(self, value_coder, window_coder, payload):
super().__init__(value_coder, TimestampCoderImpl(), window_coder)
self._timestamp, self._windows, self._pane_info = self._from_proto(
@@ -1595,6 +1623,7 @@ class LengthPrefixCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees.
Coder which prefixes the length of the encoded object in the stream."""
+
def __init__(self, value_coder):
# type: (CoderImpl) -> None
self._value_coder = value_coder
@@ -1625,6 +1654,7 @@ class ShardedKeyCoderImpl(StreamCoderImpl):
shard id byte string
encoded user key
"""
+
def __init__(self, key_coder_impl):
self._shard_id_coder_impl = BytesCoderImpl()
self._key_coder_impl = key_coder_impl
@@ -1660,6 +1690,7 @@ class TimestampPrefixingWindowCoderImpl(StreamCoderImpl):
window's max_timestamp()
encoded window using it's own coder.
"""
+
def __init__(self, window_coder_impl: CoderImpl) -> None:
self._window_coder_impl = window_coder_impl
@@ -1687,6 +1718,7 @@ def _create_opaque_window(end, encoded_window):
from apache_beam.transforms.window import BoundedWindow
class _OpaqueWindow(BoundedWindow):
+
def __init__(self, end, encoded_window):
super().__init__(end)
self.encoded_window = encoded_window
@@ -1715,6 +1747,7 @@ class TimestampPrefixingOpaqueWindowCoderImpl(StreamCoderImpl):
window's max_timestamp()
length prefixed encoded window
"""
+
def __init__(self) -> None:
pass
@@ -1770,6 +1803,7 @@ class RowColumnEncoder:
class GenericRowColumnEncoder(RowColumnEncoder):
+
def __init__(self, coder_impl, column):
self.coder_impl = coder_impl
self.column = column
@@ -1790,6 +1824,7 @@ class GenericRowColumnEncoder(RowColumnEncoder):
class RowCoderImpl(StreamCoderImpl):
"""For internal use only; no backwards-compatibility guarantees."""
+
def __init__(self, schema, components):
self.schema = schema
self.num_fields = len(self.schema.fields)
@@ -1859,8 +1894,7 @@ class RowCoderImpl(StreamCoderImpl):
RowColumnEncoder.create(
self.schema.fields[i].type.atomic_type,
self.components[i],
- columns[name]) for i,
- name in enumerate(self.field_names)
+ columns[name]) for i, name in enumerate(self.field_names)
]
def encode_batch_to_stream(self, columns: Dict[str, np.ndarray], out):
@@ -1964,6 +1998,7 @@ class RowCoderImpl(StreamCoderImpl):
class LogicalTypeCoderImpl(StreamCoderImpl):
+
def __init__(self, logical_type, representation_coder):
self.logical_type = logical_type
self.representation_coder = representation_coder.get_impl()
@@ -1982,6 +2017,7 @@ class BigIntegerCoderImpl(StreamCoderImpl):
For interoperability with Java SDK, encoding needs to match that of the Java
SDK BigIntegerCoder."""
+
def encode_to_stream(self, value, out, nested):
# type: (int, create_OutputStream, bool) -> None
if value < 0:
diff --git a/sdks/python/apache_beam/coders/coders.py b/sdks/python/apache_beam/coders/coders.py
index 0f2a4268685..54bbcb3a372 100644
--- a/sdks/python/apache_beam/coders/coders.py
+++ b/sdks/python/apache_beam/coders/coders.py
@@ -139,6 +139,7 @@ def deserialize_coder(serialized):
class Coder(object):
"""Base class for coders."""
+
def encode(self, value):
# type: (Any) -> bytes
@@ -331,6 +332,7 @@ class Coder(object):
A corresponding to_runner_api_parameter method would be expected that
returns the tuple ('beam:fn:foo', FooPayload)
"""
+
def register(fn):
cls._known_urns[urn] = parameter_type, fn
return fn
@@ -382,9 +384,8 @@ class Coder(object):
"""
setattr(
cls,
- 'to_runner_api_parameter',
- lambda self,
- unused_context: (urn, None, self._get_component_coders()))
+ 'to_runner_api_parameter', lambda self, unused_context:
+ (urn, None, self._get_component_coders()))
# pylint: disable=unused-variable
@Coder.register_urn(urn, None)
@@ -403,6 +404,7 @@ def _pickle_from_runner_api_parameter(payload, components, context):
class StrUtf8Coder(Coder):
"""A coder used for reading and writing strings as UTF-8."""
+
def encode(self, value):
return value.encode('utf-8')
@@ -422,6 +424,7 @@ Coder.register_structured_urn(common_urns.coders.STRING_UTF8.urn, StrUtf8Coder)
class ToBytesCoder(Coder):
"""A default string coder used if no sink coder is specified."""
+
def encode(self, value):
return value if isinstance(value, bytes) else str(value).encode('utf-8')
@@ -444,6 +447,7 @@ class FastCoder(Coder):
this class inverts that by defining encode() and decode() in terms of
_create_impl().
"""
+
def encode(self, value):
"""Encodes the given object into a byte string."""
return self.get_impl().encode(value)
@@ -461,6 +465,7 @@ class FastCoder(Coder):
class BytesCoder(FastCoder):
"""Byte string coder."""
+
def _create_impl(self):
return coder_impl.BytesCoderImpl()
@@ -482,6 +487,7 @@ Coder.register_structured_urn(common_urns.coders.BYTES.urn, BytesCoder)
class BooleanCoder(FastCoder):
+
def _create_impl(self):
return coder_impl.BooleanCoderImpl()
@@ -503,6 +509,7 @@ Coder.register_structured_urn(common_urns.coders.BOOL.urn, BooleanCoder)
class MapCoder(FastCoder):
+
def __init__(self, key_coder, value_coder):
# type: (Coder, Coder) -> None
self._key_coder = key_coder
@@ -548,6 +555,7 @@ class MapCoder(FastCoder):
# This is a separate class from MapCoder as the former is a standard coder with
# no way to carry the is_deterministic bit.
class DeterministicMapCoder(FastCoder):
+
def __init__(self, key_coder, value_coder):
# type: (Coder, Coder) -> None
assert key_coder.is_deterministic()
@@ -576,6 +584,7 @@ class DeterministicMapCoder(FastCoder):
class NullableCoder(FastCoder):
+
def __init__(self, value_coder):
# type: (Coder) -> None
self._value_coder = value_coder
@@ -630,6 +639,7 @@ Coder.register_structured_urn(common_urns.coders.NULLABLE.urn, NullableCoder)
class VarIntCoder(FastCoder):
"""Variable-length integer coder."""
+
def _create_impl(self):
return coder_impl.VarIntCoderImpl()
@@ -652,6 +662,7 @@ Coder.register_structured_urn(common_urns.coders.VARINT.urn, VarIntCoder)
class BigEndianShortCoder(FastCoder):
"""A coder used for big-endian int16 values."""
+
def _create_impl(self):
return coder_impl.BigEndianShortCoderImpl()
@@ -671,6 +682,7 @@ class BigEndianShortCoder(FastCoder):
class SinglePrecisionFloatCoder(FastCoder):
"""A coder used for single-precision floating-point values."""
+
def _create_impl(self):
return coder_impl.SinglePrecisionFloatCoderImpl()
@@ -696,6 +708,7 @@ class FloatCoder(FastCoder):
:class:`SinglePrecisionFloatCoder` for a single-precision version of this
coder.
"""
+
def _create_impl(self):
return coder_impl.FloatCoderImpl()
@@ -718,6 +731,7 @@ Coder.register_structured_urn(common_urns.coders.DOUBLE.urn, FloatCoder)
class TimestampCoder(FastCoder):
"""A coder used for timeutil.Timestamp values."""
+
def _create_impl(self):
return coder_impl.TimestampCoderImpl()
@@ -736,6 +750,7 @@ class _TimerCoder(FastCoder):
"""A coder used for timer values.
For internal use."""
+
def __init__(self, key_coder, window_coder):
# type: (Coder, Coder) -> None
self._key_coder = key_coder
@@ -769,6 +784,7 @@ Coder.register_structured_urn(common_urns.coders.TIMER.urn, _TimerCoder)
class SingletonCoder(FastCoder):
"""A coder that always encodes exactly one value."""
+
def __init__(self, value):
self._value = value
@@ -806,6 +822,7 @@ def maybe_dill_loads(o):
class _PickleCoderBase(FastCoder):
"""Base class for pickling coders."""
+
def is_deterministic(self):
# type: () -> bool
# Note that the default coder, the PickleCoder, is not deterministic (for
@@ -836,6 +853,7 @@ class _PickleCoderBase(FastCoder):
class _MemoizingPickleCoder(_PickleCoderBase):
"""Coder using Python's pickle functionality with memoization."""
+
def __init__(self, cache_size=16):
super().__init__()
self.cache_size = cache_size
@@ -863,6 +881,7 @@ class _MemoizingPickleCoder(_PickleCoderBase):
class PickleCoder(_PickleCoderBase):
"""Coder using Python's pickle functionality."""
+
def _create_impl(self):
dumps = pickle.dumps
protocol = pickle.HIGHEST_PROTOCOL
@@ -878,12 +897,14 @@ class PickleCoder(_PickleCoderBase):
class DillCoder(_PickleCoderBase):
"""Coder using dill's pickle functionality."""
+
def _create_impl(self):
return coder_impl.CallbackCoderImpl(maybe_dill_dumps, maybe_dill_loads)
class DeterministicFastPrimitivesCoder(FastCoder):
"""Throws runtime errors when encoding non-deterministic values."""
+
def __init__(self, coder, step_label):
self._underlying_coder = coder
self._step_label = step_label
@@ -916,6 +937,7 @@ class FastPrimitivesCoder(FastCoder):
For unknown types, falls back to another coder (e.g. PickleCoder).
"""
+
def __init__(self, fallback_coder=PickleCoder()):
# type: (Coder) -> None
self._fallback_coder = fallback_coder
@@ -962,6 +984,7 @@ class FakeDeterministicFastPrimitivesCoder(FastPrimitivesCoder):
This can be registered as a fallback coder to go back to the behavior before
deterministic encoding was enforced (BEAM-11719).
"""
+
def is_deterministic(self):
return True
@@ -1013,6 +1036,7 @@ class ProtoCoder(FastCoder):
any protobuf Message object.
"""
+
def __init__(self, proto_message_type):
# type: (Type[google.protobuf.message.Message]) -> None
self.proto_message_type = proto_message_type
@@ -1065,6 +1089,7 @@ class DeterministicProtoCoder(ProtoCoder):
version of the protoc compiler what was used to generate the protobuf
messages.
"""
+
def _create_impl(self):
return coder_impl.DeterministicProtoCoderImpl(self.proto_message_type)
@@ -1082,6 +1107,7 @@ class ProtoPlusCoder(FastCoder):
ProtoPlusCoder is registered in the global CoderRegistry as the default coder
for any proto.Message object.
"""
+
def __init__(self, proto_plus_message_type):
# type: (Type[proto.Message]) -> None
self.proto_plus_message_type = proto_plus_message_type
@@ -1117,6 +1143,7 @@ AVRO_GENERIC_CODER_URN = "beam:coder:avro:generic:v1"
class AvroGenericCoder(FastCoder):
"""A coder used for AvroRecord values."""
+
def __init__(self, schema):
self.schema = schema
@@ -1148,6 +1175,7 @@ class AvroGenericCoder(FastCoder):
class TupleCoder(FastCoder):
"""Coder of tuple objects."""
+
def __init__(self, components):
# type: (Iterable[Coder]) -> None
self._coders = tuple(components)
@@ -1224,6 +1252,7 @@ class TupleCoder(FastCoder):
class TupleSequenceCoder(FastCoder):
"""Coder of homogeneous tuple objects."""
+
def __init__(self, elem_coder):
# type: (Coder) -> None
self._elem_coder = elem_coder
@@ -1267,6 +1296,7 @@ class TupleSequenceCoder(FastCoder):
class ListLikeCoder(FastCoder):
"""Coder of iterables of homogeneous objects."""
+
def __init__(self, elem_coder):
# type: (Coder) -> None
self._elem_coder = elem_coder
@@ -1310,6 +1340,7 @@ class ListLikeCoder(FastCoder):
class IterableCoder(ListLikeCoder):
"""Coder of iterables of homogeneous objects."""
+
def to_type_hint(self):
return typehints.Iterable[self._elem_coder.to_type_hint()]
@@ -1319,6 +1350,7 @@ Coder.register_structured_urn(common_urns.coders.ITERABLE.urn, IterableCoder)
class ListCoder(ListLikeCoder):
"""Coder of Python lists."""
+
def to_type_hint(self):
return typehints.List[self._elem_coder.to_type_hint()]
@@ -1328,6 +1360,7 @@ class ListCoder(ListLikeCoder):
class GlobalWindowCoder(SingletonCoder):
"""Coder for global windows."""
+
def __init__(self):
from apache_beam.transforms import window
super().__init__(window.GlobalWindow())
@@ -1339,6 +1372,7 @@ Coder.register_structured_urn(
class IntervalWindowCoder(FastCoder):
"""Coder for an window defined by a start timestamp and a duration."""
+
def _create_impl(self):
return coder_impl.IntervalWindowCoderImpl()
@@ -1358,6 +1392,7 @@ Coder.register_structured_urn(
class _OrderedUnionCoder(FastCoder):
+
def __init__(
self, *coder_types: Tuple[type, Coder], fallback_coder: Optional[Coder]):
self._coder_types = coder_types
@@ -1390,6 +1425,7 @@ class _OrderedUnionCoder(FastCoder):
class WindowedValueCoder(FastCoder):
"""Coder for windowed values."""
+
def __init__(self, wrapped_value_coder, window_coder=None):
# type: (Coder, Optional[Coder]) -> None
if not window_coder:
@@ -1452,6 +1488,7 @@ Coder.register_structured_urn(
class ParamWindowedValueCoder(WindowedValueCoder):
"""A coder used for parameterized windowed values."""
+
def __init__(self, payload, components):
super().__init__(components[0], components[1])
self.payload = payload
@@ -1494,6 +1531,7 @@ class LengthPrefixCoder(FastCoder):
"""For internal use only; no backwards-compatibility guarantees.
Coder which prefixes the length of the encoded object in the stream."""
+
def __init__(self, value_coder):
# type: (Coder) -> None
self._value_coder = value_coder
@@ -1595,6 +1633,7 @@ class StateBackedIterableCoder(FastCoder):
class ShardedKeyCoder(FastCoder):
"""A coder for sharded key."""
+
def __init__(self, key_coder):
# type: (Coder) -> None
self._key_coder = key_coder
@@ -1644,6 +1683,7 @@ class TimestampPrefixingWindowCoder(FastCoder):
Coder which prefixes the max timestamp of arbitrary window to its encoded
form."""
+
def __init__(self, window_coder: Coder) -> None:
self._window_coder = window_coder
@@ -1679,6 +1719,7 @@ class TimestampPrefixingOpaqueWindowCoder(FastCoder):
"""For internal use only; no backwards-compatibility guarantees.
Coder which decodes windows as bytes."""
+
def __init__(self) -> None:
pass
@@ -1704,6 +1745,7 @@ Coder.register_structured_urn(
class BigIntegerCoder(FastCoder):
+
def _create_impl(self):
return coder_impl.BigIntegerCoderImpl()
@@ -1722,6 +1764,7 @@ class BigIntegerCoder(FastCoder):
class DecimalCoder(FastCoder):
+
def _create_impl(self):
return coder_impl.DecimalCoderImpl()
diff --git a/sdks/python/apache_beam/coders/coders_property_based_test.py b/sdks/python/apache_beam/coders/coders_property_based_test.py
index 9279fc31c09..c80418f0f3e 100644
--- a/sdks/python/apache_beam/coders/coders_property_based_test.py
+++ b/sdks/python/apache_beam/coders/coders_property_based_test.py
@@ -91,6 +91,7 @@ TYPES_UNSUPPORTED_BY_ROW_CODER = {np.int8, np.int16}
class TypesAreAllTested(unittest.TestCase):
+
def test_all_types_are_tested(self):
# Verify that all types among Beam's defined types are being tested
self.assertEqual(
@@ -100,6 +101,7 @@ class TypesAreAllTested(unittest.TestCase):
class ProperyTestingCoders(unittest.TestCase):
+
@given(st.text())
def test_string_coder(self, txt: str):
coder = StrUtf8Coder()
@@ -144,9 +146,7 @@ class ProperyTestingCoders(unittest.TestCase):
row = RowType(
**{
name: data.draw(SCHEMA_TYPES_TO_STRATEGY[type_])
- for name,
- type_,
- nullable in schema
+ for name, type_, nullable in schema
})
coder = RowCoder(typing_to_runner_api(RowType).row_type.schema)
diff --git a/sdks/python/apache_beam/coders/coders_test.py b/sdks/python/apache_beam/coders/coders_test.py
index 5e5debca36e..c4f6f239ee9 100644
--- a/sdks/python/apache_beam/coders/coders_test.py
+++ b/sdks/python/apache_beam/coders/coders_test.py
@@ -34,6 +34,7 @@ from apache_beam.testing.test_pipeline import TestPipeline
class PickleCoderTest(unittest.TestCase):
+
def test_basics(self):
v = ('a' * 10, 'b' * 90)
pickler = coders.PickleCoder()
@@ -52,6 +53,7 @@ class PickleCoderTest(unittest.TestCase):
class CodersTest(unittest.TestCase):
+
def test_str_utf8_coder(self):
real_coder = coders_registry.get_coder(bytes)
expected_coder = coders.BytesCoder()
@@ -75,6 +77,7 @@ class CodersTest(unittest.TestCase):
# TODO(https://github.com/apache/beam/issues/22319): The proto file should be
# placed in a common directory that can be shared between java and python.
class ProtoCoderTest(unittest.TestCase):
+
def test_proto_coder(self):
ma = test_message.MessageA()
mb = ma.field2.add()
@@ -106,6 +109,7 @@ class ProtoCoderTest(unittest.TestCase):
class DeterministicProtoCoderTest(unittest.TestCase):
+
def test_deterministic_proto_coder(self):
ma = test_message.MessageA()
mb = ma.field2.add()
@@ -147,6 +151,7 @@ class ProtoPlusMessageWithMap(proto.Message):
class ProtoPlusCoderTest(unittest.TestCase):
+
def test_proto_plus_coder(self):
ma = ProtoPlusMessageA()
ma.field2 = [ProtoPlusMessageB(field1=True)]
@@ -195,6 +200,7 @@ coders_registry.register_coder(AvroTestRecord, AvroTestCoder)
class AvroCoderTest(unittest.TestCase):
+
def test_avro_record_coder(self):
real_coder = coders_registry.get_coder(AvroTestRecord)
expected_coder = AvroTestCoder()
@@ -219,6 +225,7 @@ class AvroCoderTest(unittest.TestCase):
class DummyClass(object):
"""A class with no registered coder."""
+
def __init__(self):
pass
@@ -232,6 +239,7 @@ class DummyClass(object):
class FallbackCoderTest(unittest.TestCase):
+
def test_default_fallback_path(self):
"""Test fallback path picks a matching coder if no coder is registered."""
@@ -243,6 +251,7 @@ class FallbackCoderTest(unittest.TestCase):
class NullableCoderTest(unittest.TestCase):
+
def test_determinism(self):
deterministic = coders_registry.get_coder(typehints.Optional[int])
deterministic.as_deterministic_coder('label')
@@ -257,12 +266,14 @@ class NullableCoderTest(unittest.TestCase):
class LengthPrefixCoderTest(unittest.TestCase):
+
def test_to_type_hint(self):
coder = coders.LengthPrefixCoder(coders.BytesCoder())
assert coder.to_type_hint() is bytes
class NumpyIntAsKeyTest(unittest.TestCase):
+
def test_numpy_int(self):
# this type is not supported as the key
import numpy as np
diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py
index f3381cdb1d6..5ce45b1aeab 100644
--- a/sdks/python/apache_beam/coders/coders_test_common.py
+++ b/sdks/python/apache_beam/coders/coders_test_common.py
@@ -69,6 +69,7 @@ MyFlag = enum.Flag('MyFlag', 'F1 F2 F3') # pylint: disable=too-many-function-ar
class DefinesGetState:
+
def __init__(self, value):
self.value = value
@@ -80,12 +81,14 @@ class DefinesGetState:
class DefinesGetAndSetState(DefinesGetState):
+
def __setstate__(self, value):
self.value = value
# Defined out of line for picklability.
class CustomCoder(coders.Coder):
+
def encode(self, x):
return str(x + 1).encode('utf-8')
@@ -163,7 +166,7 @@ class CodersTest(unittest.TestCase):
coders.BigEndianShortCoder,
coders.SinglePrecisionFloatCoder,
coders.ToBytesCoder,
- coders.BigIntegerCoder, # tested in DecimalCoder
+ coders.BigIntegerCoder, # tested in DecimalCoder
coders.TimestampPrefixingOpaqueWindowCoder,
])
cls.seen_nested -= set(
@@ -434,6 +437,7 @@ class CodersTest(unittest.TestCase):
self._test_iterable_coder_of_unknown_length(80000)
def _test_iterable_coder_of_unknown_length(self, count):
+
def iter_generator(count):
for i in range(count):
yield i
@@ -614,7 +618,9 @@ class CodersTest(unittest.TestCase):
coders.TupleCoder((coder, coder)), (b'', b'a'), (b'bc', b'def'))
def test_nested_observables(self):
+
class FakeObservableIterator(observable.ObservableMixin):
+
def __iter__(self):
return iter([1, 2, 3])
@@ -672,9 +678,14 @@ class CodersTest(unittest.TestCase):
def test_map_coder(self):
values = [
- {1: "one", 300: "three hundred"}, # force yapf to be nice
+ {
+ 1: "one", 300: "three hundred"
+ }, # force yapf to be nice
{},
- {i: str(i) for i in range(5000)}
+ {
+ i: str(i)
+ for i in range(5000)
+ }
]
map_coder = coders.MapCoder(coders.VarIntCoder(), coders.StrUtf8Coder())
self.check_coder(map_coder, *values)
diff --git a/sdks/python/apache_beam/coders/fast_coders_test.py b/sdks/python/apache_beam/coders/fast_coders_test.py
index fa8643c2a38..8fc7e4a3bdb 100644
--- a/sdks/python/apache_beam/coders/fast_coders_test.py
+++ b/sdks/python/apache_beam/coders/fast_coders_test.py
@@ -27,6 +27,7 @@ from apache_beam.tools import utils
class FastCoders(unittest.TestCase):
+
def test_using_fast_impl(self):
try:
utils.check_compiled('apache_beam.coders.coder_impl')
diff --git a/sdks/python/apache_beam/coders/observable.py b/sdks/python/apache_beam/coders/observable.py
index 8cdc3227ec8..28ad58876de 100644
--- a/sdks/python/apache_beam/coders/observable.py
+++ b/sdks/python/apache_beam/coders/observable.py
@@ -29,6 +29,7 @@ class ObservableMixin(object):
Subclasses need to call self.notify_observers with any object yielded.
"""
+
def __init__(self):
self.observers = []
diff --git a/sdks/python/apache_beam/coders/observable_test.py b/sdks/python/apache_beam/coders/observable_test.py
index df4e7ef0940..b1a59a19288 100644
--- a/sdks/python/apache_beam/coders/observable_test.py
+++ b/sdks/python/apache_beam/coders/observable_test.py
@@ -37,7 +37,9 @@ class ObservableMixinTest(unittest.TestCase):
self.observed_keys.append(key)
def test_observable(self):
+
class Watched(observable.ObservableMixin):
+
def __iter__(self):
for i in (1, 4, 3):
self.notify_observers(i, key='a%d' % i)
diff --git a/sdks/python/apache_beam/coders/row_coder.py b/sdks/python/apache_beam/coders/row_coder.py
index e93abbc887f..3b39c00b7f0 100644
--- a/sdks/python/apache_beam/coders/row_coder.py
+++ b/sdks/python/apache_beam/coders/row_coder.py
@@ -51,6 +51,7 @@ class RowCoder(FastCoder):
Implements the beam:coder:row:v1 standard coder spec.
"""
+
def __init__(self, schema, force_deterministic=False):
"""Initializes a :class:`RowCoder`.
@@ -189,6 +190,7 @@ def _nonnull_coder_from_type(field_type):
class LogicalTypeCoder(FastCoder):
+
def __init__(self, logical_type, representation_coder):
self.logical_type = logical_type
self.representation_coder = representation_coder
diff --git a/sdks/python/apache_beam/coders/row_coder_test.py b/sdks/python/apache_beam/coders/row_coder_test.py
index 6ac982835cb..440a4182f97 100644
--- a/sdks/python/apache_beam/coders/row_coder_test.py
+++ b/sdks/python/apache_beam/coders/row_coder_test.py
@@ -423,8 +423,7 @@ class RowCoderTest(unittest.TestCase):
for size in [len(self.PEOPLE) - 1, len(self.PEOPLE), len(self.PEOPLE) + 1]:
dest = {
field: np.ndarray((size, ), dtype=a.dtype)
- for field,
- a in columnar.items()
+ for field, a in columnar.items()
}
n = min(size, len(self.PEOPLE))
self.assertEqual(
diff --git a/sdks/python/apache_beam/coders/slow_coders_test.py b/sdks/python/apache_beam/coders/slow_coders_test.py
index 7915116a19a..ab325edd5e2 100644
--- a/sdks/python/apache_beam/coders/slow_coders_test.py
+++ b/sdks/python/apache_beam/coders/slow_coders_test.py
@@ -29,6 +29,7 @@ from apache_beam.coders.coders_test_common import *
'Remove non-cython tests.'
'https://github.com/apache/beam/issues/28307')
class SlowCoders(unittest.TestCase):
+
def test_using_slow_impl(self):
try:
# pylint: disable=wrong-import-position
diff --git a/sdks/python/apache_beam/coders/slow_stream.py b/sdks/python/apache_beam/coders/slow_stream.py
index b08ad8e9a37..f4cfb5b8273 100644
--- a/sdks/python/apache_beam/coders/slow_stream.py
+++ b/sdks/python/apache_beam/coders/slow_stream.py
@@ -29,6 +29,7 @@ class OutputStream(object):
"""For internal use only; no backwards-compatibility guarantees.
A pure Python implementation of stream.OutputStream."""
+
def __init__(self):
self.data: List[bytes] = []
self.byte_count = 0
@@ -91,6 +92,7 @@ class ByteCountingOutputStream(OutputStream):
"""For internal use only; no backwards-compatibility guarantees.
A pure Python implementation of stream.ByteCountingOutputStream."""
+
def __init__(self):
# Note that we don't actually use any of the data initialized by our super.
super().__init__()
@@ -119,6 +121,7 @@ class InputStream(object):
"""For internal use only; no backwards-compatibility guarantees.
A pure Python implementation of stream.InputStream."""
+
def __init__(self, data: bytes) -> None:
self.data = data
self.pos = 0
diff --git a/sdks/python/apache_beam/coders/standard_coders_test.py b/sdks/python/apache_beam/coders/standard_coders_test.py
index 47df0116f2c..59b7024188a 100644
--- a/sdks/python/apache_beam/coders/standard_coders_test.py
+++ b/sdks/python/apache_beam/coders/standard_coders_test.py
@@ -77,6 +77,7 @@ def parse_float(s):
def value_parser_from_schema(schema):
+
def attribute_parser_from_type(type_):
parser = nonnull_attribute_parser_from_type(type_)
if type_.nullable:
@@ -139,65 +140,48 @@ class StandardCodersTest(unittest.TestCase):
'beam:coder:bool:v1': lambda x: x,
'beam:coder:string_utf8:v1': lambda x: x,
'beam:coder:varint:v1': lambda x: x,
- 'beam:coder:kv:v1': lambda x,
- key_parser,
- value_parser: (key_parser(x['key']), value_parser(x['value'])),
+ 'beam:coder:kv:v1': lambda x, key_parser, value_parser:
+ (key_parser(x['key']), value_parser(x['value'])),
'beam:coder:interval_window:v1': lambda x: IntervalWindow(
- start=Timestamp(micros=(x['end'] - x['span']) * 1000),
- end=Timestamp(micros=x['end'] * 1000)),
- 'beam:coder:iterable:v1': lambda x,
- parser: list(map(parser, x)),
- 'beam:coder:state_backed_iterable:v1': lambda x,
- parser: list(map(parser, x)),
+ start=Timestamp(micros=(x['end'] - x['span']) * 1000), end=Timestamp(
+ micros=x['end'] * 1000)),
+ 'beam:coder:iterable:v1': lambda x, parser: list(map(parser, x)),
+ 'beam:coder:state_backed_iterable:v1': lambda x, parser: list(
+ map(parser, x)),
'beam:coder:global_window:v1': lambda x: window.GlobalWindow(),
- 'beam:coder:windowed_value:v1': lambda x,
- value_parser,
- window_parser: windowed_value.create(
- value_parser(x['value']),
- x['timestamp'] * 1000,
- tuple(window_parser(w) for w in x['windows'])),
- 'beam:coder:param_windowed_value:v1': lambda x,
- value_parser,
+ 'beam:coder:windowed_value:v1': lambda x, value_parser, window_parser:
+ windowed_value.create(
+ value_parser(x['value']), x['timestamp'] * 1000, tuple(
+ window_parser(w) for w in x['windows'])),
+ 'beam:coder:param_windowed_value:v1': lambda x, value_parser,
window_parser: windowed_value.create(
- value_parser(x['value']),
- x['timestamp'] * 1000,
- tuple(window_parser(w) for w in x['windows']),
- PaneInfo(
- x['pane']['is_first'],
- x['pane']['is_last'],
- PaneInfoTiming.from_string(x['pane']['timing']),
- x['pane']['index'],
- x['pane']['on_time_index'])),
- 'beam:coder:timer:v1': lambda x,
- value_parser,
- window_parser: userstate.Timer(
- user_key=value_parser(x['userKey']),
- dynamic_timer_tag=x['dynamicTimerTag'],
- clear_bit=x['clearBit'],
- windows=tuple(window_parser(w) for w in x['windows']),
- fire_timestamp=None,
- hold_timestamp=None,
- paneinfo=None) if x['clearBit'] else userstate.Timer(
- user_key=value_parser(x['userKey']),
- dynamic_timer_tag=x['dynamicTimerTag'],
- clear_bit=x['clearBit'],
- fire_timestamp=Timestamp(micros=x['fireTimestamp'] * 1000),
- hold_timestamp=Timestamp(micros=x['holdTimestamp'] * 1000),
- windows=tuple(window_parser(w) for w in x['windows']),
- paneinfo=PaneInfo(
- x['pane']['is_first'],
- x['pane']['is_last'],
- PaneInfoTiming.from_string(x['pane']['timing']),
- x['pane']['index'],
- x['pane']['on_time_index'])),
+ value_parser(x['value']), x['timestamp'] * 1000, tuple(
+ window_parser(w) for w in x['windows']), PaneInfo(
+ x['pane']['is_first'], x['pane']['is_last'], PaneInfoTiming.
+ from_string(x['pane']['timing']), x['pane']['index'], x[
+ 'pane']['on_time_index'])),
+ 'beam:coder:timer:v1': lambda x, value_parser, window_parser: userstate.
+ Timer(
+ user_key=value_parser(x['userKey']), dynamic_timer_tag=x[
+ 'dynamicTimerTag'], clear_bit=x['clearBit'], windows=tuple(
+ window_parser(w) for w in x['windows']), fire_timestamp=None,
+ hold_timestamp=None, paneinfo=None)
+ if x['clearBit'] else userstate.Timer(
+ user_key=value_parser(x['userKey']), dynamic_timer_tag=x[
+ 'dynamicTimerTag'], clear_bit=x['clearBit'], fire_timestamp=
+ Timestamp(micros=x['fireTimestamp'] * 1000), hold_timestamp=Timestamp(
+ micros=x['holdTimestamp'] * 1000), windows=tuple(
+ window_parser(w) for w in x['windows']), paneinfo=PaneInfo(
+ x['pane']['is_first'], x['pane']['is_last'],
+ PaneInfoTiming.from_string(x['pane']['timing']), x[
+ 'pane']['index'], x['pane']['on_time_index'])),
'beam:coder:double:v1': parse_float,
- 'beam:coder:sharded_key:v1': lambda x,
- value_parser: ShardedKey(
+ 'beam:coder:sharded_key:v1': lambda x, value_parser: ShardedKey(
key=value_parser(x['key']), shard_id=x['shardId'].encode('utf-8')),
- 'beam:coder:custom_window:v1': lambda x,
- window_parser: window_parser(x['window']),
- 'beam:coder:nullable:v1': lambda x,
- value_parser: x.encode('utf-8') if x else None
+ 'beam:coder:custom_window:v1': lambda x, window_parser: window_parser(
+ x['window']),
+ 'beam:coder:nullable:v1': lambda x, value_parser: x.encode('utf-8')
+ if x else None
}
def test_standard_coders(self):
@@ -206,6 +190,7 @@ class StandardCodersTest(unittest.TestCase):
self._run_standard_coder(name, spec)
def _run_standard_coder(self, name, spec):
+
def assert_equal(actual, expected):
"""Handle nan values which self.assertEqual fails on."""
if (isinstance(actual, float) and isinstance(expected, float) and
diff --git a/sdks/python/apache_beam/coders/typecoders.py b/sdks/python/apache_beam/coders/typecoders.py
index 1667cb7a916..fdc66d5d0f6 100644
--- a/sdks/python/apache_beam/coders/typecoders.py
+++ b/sdks/python/apache_beam/coders/typecoders.py
@@ -79,6 +79,7 @@ __all__ = ['registry']
class CoderRegistry(object):
"""A coder registry for typehint/coder associations."""
+
def __init__(self, fallback_coder=None):
self._coders: Dict[Any, Type[coders.Coder]] = {}
self.custom_types: List[Any] = []
@@ -188,6 +189,7 @@ class FirstOf(object):
"""For internal use only; no backwards-compatibility guarantees.
A class used to get the first matching coder from a list of coders."""
+
def __init__(self, coders: Iterable[Type[coders.Coder]]) -> None:
self._coders = coders
diff --git a/sdks/python/apache_beam/coders/typecoders_test.py b/sdks/python/apache_beam/coders/typecoders_test.py
index 3adc8255409..03e52a0802c 100644
--- a/sdks/python/apache_beam/coders/typecoders_test.py
+++ b/sdks/python/apache_beam/coders/typecoders_test.py
@@ -27,6 +27,7 @@ from apache_beam.typehints import typehints
class CustomClass(object):
+
def __init__(self, n):
self.number = n
@@ -38,6 +39,7 @@ class CustomClass(object):
class CustomCoder(coders.Coder):
+
def encode(self, value):
return str(value.number).encode('ASCII')
@@ -52,6 +54,7 @@ class CustomCoder(coders.Coder):
class TypeCodersTest(unittest.TestCase):
+
def test_register_non_type_coder(self):
coder = CustomCoder()
with self.assertRaisesRegex(
diff --git a/sdks/python/apache_beam/dataframe/convert.py b/sdks/python/apache_beam/dataframe/convert.py
index c5a0d1025c6..52b4900cf79 100644
--- a/sdks/python/apache_beam/dataframe/convert.py
+++ b/sdks/python/apache_beam/dataframe/convert.py
@@ -92,12 +92,14 @@ UNBATCHED_CACHE: 'weakref.WeakValueDictionary[str, pvalue.PCollection]' = (
class RowsToDataFrameFn(beam.DoFn):
+
@beam.DoFn.yields_elements
def process_batch(self, batch: pd.DataFrame) -> Iterable[pd.DataFrame]:
yield batch
class ElementsToSeriesFn(beam.DoFn):
+
@beam.DoFn.yields_elements
def process_batch(self, batch: pd.Series) -> Iterable[pd.Series]:
yield batch
@@ -136,6 +138,7 @@ def _make_unbatched_pcoll(
class DataFrameToRowsFn(beam.DoFn):
+
def __init__(self, proxy, include_indexes):
self._proxy = proxy
self._include_indexes = include_indexes
@@ -150,6 +153,7 @@ class DataFrameToRowsFn(beam.DoFn):
class SeriesToElementsFn(beam.DoFn):
+
def __init__(self, proxy):
self._proxy = proxy
@@ -257,14 +261,14 @@ def to_pcollection(
{ix: df._expr
for (ix, df) in enumerate(new_dataframes)})
- TO_PCOLLECTION_CACHE.update(
- {new_dataframes[ix]._expr._id: pc
- for ix, pc in new_results.items()})
+ TO_PCOLLECTION_CACHE.update({
+ new_dataframes[ix]._expr._id: pc
+ for ix, pc in new_results.items()
+ })
raw_results = {
ix: TO_PCOLLECTION_CACHE[df._expr._id]
- for ix,
- df in enumerate(dataframes)
+ for ix, df in enumerate(dataframes)
}
if yield_elements == "schemas":
diff --git a/sdks/python/apache_beam/dataframe/convert_test.py b/sdks/python/apache_beam/dataframe/convert_test.py
index b00ce0e51fa..3c78c8a4343 100644
--- a/sdks/python/apache_beam/dataframe/convert_test.py
+++ b/sdks/python/apache_beam/dataframe/convert_test.py
@@ -25,6 +25,7 @@ from apache_beam.testing.util import equal_to
def equal_to_unordered_series(expected):
+
def check(actual):
actual = pd.concat(actual)
if sorted(expected) != sorted(actual):
@@ -34,6 +35,7 @@ def equal_to_unordered_series(expected):
class ConvertTest(unittest.TestCase):
+
def test_convert_yield_pandas(self):
with beam.Pipeline() as p:
a = pd.Series([1, 2, 3])
@@ -185,7 +187,9 @@ class ConvertTest(unittest.TestCase):
logging.disable(logging.NOTSET)
def test_auto_convert(self):
+
class MySchemaTransform(beam.PTransform):
+
def expand(self, pcoll):
return pcoll | beam.Map(
lambda x: beam.Row(
diff --git a/sdks/python/apache_beam/dataframe/doctests.py b/sdks/python/apache_beam/dataframe/doctests.py
index 57ee8009ba4..52b72efde9b 100644
--- a/sdks/python/apache_beam/dataframe/doctests.py
+++ b/sdks/python/apache_beam/dataframe/doctests.py
@@ -60,6 +60,7 @@ from apache_beam.dataframe.frame_base import DeferredBase
class FakePandasObject(object):
"""A stand-in for the wrapped pandas objects.
"""
+
def __init__(self, pandas_obj, test_env):
self._pandas_obj = pandas_obj
self._test_env = test_env
@@ -94,6 +95,7 @@ class TestEnvironment(object):
These classes are patched to be able to recognize and retrieve inputs
and results, stored in `self._inputs` and `self._all_frames` respectively.
"""
+
def __init__(self):
self._inputs = {}
self._all_frames = {}
@@ -157,6 +159,7 @@ class _InMemoryResultRecorder(object):
del self._ALL_RESULTS[self._id]
def record_fn(self, name):
+
def record(value):
self._ALL_RESULTS[self._id][name].append(value)
@@ -173,6 +176,7 @@ NOT_IMPLEMENTED = 'NotImplementedError'
class _DeferrredDataframeOutputChecker(doctest.OutputChecker):
"""Validates output by replacing DeferredBase[...] with computed values.
"""
+
def __init__(self, env, use_beam):
self._env = env
if use_beam:
@@ -188,8 +192,7 @@ class _DeferrredDataframeOutputChecker(doctest.OutputChecker):
session = expressions.PartitioningSession(self._env._inputs)
return {
name: session.evaluate(frame._expr)
- for name,
- frame in to_compute.items()
+ for name, frame in to_compute.items()
}
def compute_using_beam(self, to_compute):
@@ -198,13 +201,13 @@ class _DeferrredDataframeOutputChecker(doctest.OutputChecker):
input_pcolls = {
placeholder: p
| 'Create%s' % placeholder >> beam.Create([input[::2], input[1::2]])
- for placeholder,
- input in self._env._inputs.items()
+ for placeholder, input in self._env._inputs.items()
}
output_pcolls = (
- input_pcolls | transforms._DataframeExpressionsTransform(
- {name: frame._expr
- for name, frame in to_compute.items()}))
+ input_pcolls | transforms._DataframeExpressionsTransform({
+ name: frame._expr
+ for name, frame in to_compute.items()
+ }))
for name, output_pcoll in output_pcolls.items():
_ = output_pcoll | 'Record%s' % name >> beam.FlatMap(
recorder.record_fn(name))
@@ -341,6 +344,7 @@ class BeamDataframeDoctestRunner(doctest.DocTestRunner):
"""A Doctest runner suitable for replacing the `pd` module with one backed
by beam.
"""
+
def __init__(
self,
env,
@@ -359,18 +363,15 @@ class BeamDataframeDoctestRunner(doctest.DocTestRunner):
self._wont_implement_ok = {
test: [to_callable(cond) for cond in examples]
- for test,
- examples in (wont_implement_ok or {}).items()
+ for test, examples in (wont_implement_ok or {}).items()
}
self._not_implemented_ok = {
test: [to_callable(cond) for cond in examples]
- for test,
- examples in (not_implemented_ok or {}).items()
+ for test, examples in (not_implemented_ok or {}).items()
}
self._skip = {
test: [to_callable(cond) for cond in examples]
- for test,
- examples in (skip or {}).items()
+ for test, examples in (skip or {}).items()
}
super().__init__(
checker=_DeferrredDataframeOutputChecker(self._test_env, use_beam),
@@ -420,6 +421,7 @@ class BeamDataframeDoctestRunner(doctest.DocTestRunner):
return result
def report_success(self, out, test, example, got):
+
def extract_concise_reason(got, expected_exc):
m = re.search(r"Implement(?:ed)?Error:\s+(.*)\n$", got)
if m:
@@ -463,6 +465,7 @@ class AugmentedTestResults(doctest.TestResults):
class Summary(object):
+
def __init__(self, failures=0, tries=0, skipped=0, error_reasons=None):
self.failures = failures
self.tries = tries
@@ -487,6 +490,7 @@ class Summary(object):
merged_reasons)
def summarize(self):
+
def print_partition(indent, desc, n, total):
print("%s%d %s (%.1f%%)" % (" " * indent, n, desc, n / total * 100))
@@ -535,9 +539,9 @@ def parse_rst_ipython_tests(rst, name, extraglobs=None, optionflags=None):
IMPORT_PANDAS = 'import pandas as pd'
example_srcs = []
- lines = iter([(lineno, line.rstrip()) for lineno,
- line in enumerate(rst.split('\n')) if is_example_line(line)] +
- [(None, 'END')])
+ lines = iter([(lineno, line.rstrip())
+ for lineno, line in enumerate(rst.split('\n'))
+ if is_example_line(line)] + [(None, 'END')])
# https://ipython.readthedocs.io/en/stable/sphinxext.html
lineno, line = next(lines)
@@ -610,6 +614,7 @@ def test_rst_ipython(
"""Extracts examples from an rst file and run them through pandas to get the
expected output, and then compare them against our dataframe implementation.
"""
+
def run_tests(extraglobs, optionflags, **kwargs):
# The patched one.
tests = parse_rst_ipython_tests(rst, name, extraglobs, optionflags)
@@ -690,12 +695,8 @@ def _run_patched(func, *args, **kwargs):
# Unfortunately the runner is not injectable.
original_doc_test_runner = doctest.DocTestRunner
doctest.DocTestRunner = lambda **kwargs: BeamDataframeDoctestRunner(
- env,
- use_beam=use_beam,
- wont_implement_ok=wont_implement_ok,
- not_implemented_ok=not_implemented_ok,
- skip=skip,
- **kwargs)
+ env, use_beam=use_beam, wont_implement_ok=wont_implement_ok,
+ not_implemented_ok=not_implemented_ok, skip=skip, **kwargs)
with expressions.allow_non_parallel_operations():
return func(
*args, extraglobs=extraglobs, optionflags=optionflags, **kwargs)
diff --git a/sdks/python/apache_beam/dataframe/doctests_test.py b/sdks/python/apache_beam/dataframe/doctests_test.py
index df24213c871..f966dc1f92f 100644
--- a/sdks/python/apache_beam/dataframe/doctests_test.py
+++ b/sdks/python/apache_beam/dataframe/doctests_test.py
@@ -144,6 +144,7 @@ and someting that'll fail (due to fake vs. real pandas)
class DoctestTest(unittest.TestCase):
+
def test_good(self):
result = doctests.teststring(SAMPLE_DOCTEST, report=False)
self.assertEqual(result.attempted, 3)
diff --git a/sdks/python/apache_beam/dataframe/expressions.py b/sdks/python/apache_beam/dataframe/expressions.py
index 2ef172b8dad..94a9d4d825c 100644
--- a/sdks/python/apache_beam/dataframe/expressions.py
+++ b/sdks/python/apache_beam/dataframe/expressions.py
@@ -33,6 +33,7 @@ class Session(object):
The bindings typically include required placeholders, but may be any
intermediate expression as well.
"""
+
def __init__(self, bindings=None):
self._bindings = dict(bindings or {})
@@ -60,6 +61,7 @@ class PartitioningSession(Session):
For testing only.
"""
+
def evaluate(self, expr):
import pandas as pd
import collections
@@ -205,6 +207,7 @@ class Expression(Generic[T]):
expression. However, unless the inputs are Singleton-partitioned, the
expression makes no guarantees about the partitioning of the output.
"""
+
def __init__(self, name: str, proxy: T, _id: Optional[str] = None):
self._name = name
self._proxy = proxy
@@ -250,6 +253,7 @@ class Expression(Generic[T]):
class PlaceholderExpression(Expression):
"""An expression whose value must be explicitly bound in the session."""
+
def __init__(
self,
proxy: T,
@@ -282,6 +286,7 @@ class PlaceholderExpression(Expression):
class ConstantExpression(Expression):
"""An expression whose value is known at pipeline construction time."""
+
def __init__(self, value: T, proxy: Optional[T] = None):
"""Initialize a constant expression.
@@ -314,6 +319,7 @@ class ConstantExpression(Expression):
class ComputedExpression(Expression):
"""An expression whose value must be computed at pipeline execution time."""
+
def __init__(
self,
name: str,
@@ -410,6 +416,7 @@ def allow_non_parallel_operations(allow=True):
class NonParallelOperation(Exception):
+
def __init__(self, msg):
super().__init__(self, msg)
self.msg = msg
diff --git a/sdks/python/apache_beam/dataframe/expressions_test.py b/sdks/python/apache_beam/dataframe/expressions_test.py
index 2c5c716f5be..65e428f9780 100644
--- a/sdks/python/apache_beam/dataframe/expressions_test.py
+++ b/sdks/python/apache_beam/dataframe/expressions_test.py
@@ -23,6 +23,7 @@ from apache_beam.dataframe import partitionings
class ExpressionTest(unittest.TestCase):
+
def test_placeholder_expression(self):
a = expressions.PlaceholderExpression(None)
b = expressions.PlaceholderExpression(None)
diff --git a/sdks/python/apache_beam/dataframe/frame_base.py b/sdks/python/apache_beam/dataframe/frame_base.py
index 8e206fc5e03..a700c044090 100644
--- a/sdks/python/apache_beam/dataframe/frame_base.py
+++ b/sdks/python/apache_beam/dataframe/frame_base.py
@@ -43,6 +43,7 @@ class DeferredBase(object):
@classmethod
def _register_for(cls, pandas_type):
+
def wrapper(deferred_type):
cls._pandas_type_map[pandas_type] = deferred_type
return deferred_type
@@ -91,6 +92,7 @@ class UnusableUnpickledDeferredBase(object):
Trying to use this object after unpickling is a bug and will result in an
error.
"""
+
def __init__(self, name):
self._name = name
@@ -103,6 +105,7 @@ class DeferredFrame(DeferredBase):
class _DeferredScalar(DeferredBase):
+
def apply(self, func, name=None, args=()):
if name is None:
name = func.__name__
@@ -135,6 +138,7 @@ class _DeferredScalar(DeferredBase):
def _scalar_binop(op):
+
def binop(self, other):
if not isinstance(other, DeferredBase):
return self.apply(lambda left: getattr(left, op)(other), name=op)
@@ -442,6 +446,7 @@ def not_implemented_method(op, issue='20318', base_type=None):
def _copy_and_mutate(func):
+
def wrapper(self, *args, **kwargs):
copy = self.copy()
func(copy, *args, **kwargs)
@@ -462,6 +467,7 @@ def maybe_inplace(func):
the inplace operation will refernce the updated expression.
For internal use only. No backwards compatibility guarantees."""
+
@functools.wraps(func)
def wrapper(self, inplace=False, **kwargs):
result = func(self, **kwargs)
@@ -491,6 +497,7 @@ def args_to_kwargs(base_type, removed_method=False, removed_args=None):
removed_args: If not empty, which arguments have been dropped in the
running Pandas version.
"""
+
def wrap(func):
if removed_method:
# Do no processing, let Beam function itself raise the error if called.
@@ -566,6 +573,7 @@ def with_docs_from(base_type, name=None, removed_method=False):
removed_method used in cases where a method has been removed in a later
version of Pandas.
"""
+
def wrap(func):
if removed_method:
func.__doc__ = (
@@ -659,6 +667,7 @@ def populate_defaults(base_type, removed_method=False, removed_args=None):
removed_args: If not empty, which arguments have been dropped in the
running Pandas version.
"""
+
def wrap(func):
if removed_method:
return func
@@ -750,6 +759,7 @@ class WontImplementError(NotImplementedError):
Raising this error will also prevent this doctests from being validated
when run with the beam dataframe validation doctest runner.
"""
+
def __init__(self, msg, reason=None):
if reason is not None:
if reason not in _WONT_IMPLEMENT_REASONS:
diff --git a/sdks/python/apache_beam/dataframe/frame_base_test.py b/sdks/python/apache_beam/dataframe/frame_base_test.py
index 0a73905339f..9c374314714 100644
--- a/sdks/python/apache_beam/dataframe/frame_base_test.py
+++ b/sdks/python/apache_beam/dataframe/frame_base_test.py
@@ -24,6 +24,7 @@ from apache_beam.dataframe import frames
class FrameBaseTest(unittest.TestCase):
+
def test_elementwise_func(self):
a = pd.Series([1, 2, 3])
b = pd.Series([100, 200, 300])
@@ -56,6 +57,7 @@ class FrameBaseTest(unittest.TestCase):
self.assertTrue(sub(x, y)._expr.evaluate_at(session).equals(a - b))
def test_maybe_inplace(self):
+
@frame_base.maybe_inplace
def add_one(frame):
return frame + 1
@@ -71,11 +73,14 @@ class FrameBaseTest(unittest.TestCase):
self.assertIsNot(x._expr, original_expr)
def test_args_to_kwargs(self):
+
class Base(object):
+
def func(self, a=1, b=2, c=3, *, kw_only=4):
pass
class Proxy(object):
+
@frame_base.args_to_kwargs(Base)
def func(self, **kwargs):
return kwargs
@@ -92,7 +97,9 @@ class FrameBaseTest(unittest.TestCase):
proxy.func(2, 4, 6, 8)
def test_args_to_kwargs_populates_defaults(self):
+
class Base(object):
+
def func(self, a=1, b=2, c=3):
pass
@@ -100,6 +107,7 @@ class FrameBaseTest(unittest.TestCase):
pass
class Proxy(object):
+
@frame_base.args_to_kwargs(Base)
@frame_base.populate_defaults(Base)
def func(self, a, c=1000, **kwargs):
@@ -133,11 +141,14 @@ class FrameBaseTest(unittest.TestCase):
self.assertEqual(proxy.func_removed_args(12, d=100), {'a': 12, 'd': 100})
def test_args_to_kwargs_populates_default_handles_kw_only(self):
+
class Base(object):
+
def func(self, a, b=2, c=3, *, kw_only=4):
pass
class ProxyUsesKwOnly(object):
+
@frame_base.args_to_kwargs(Base)
@frame_base.populate_defaults(Base)
def func(self, a, kw_only, **kwargs):
@@ -158,6 +169,7 @@ class FrameBaseTest(unittest.TestCase):
proxy.func(2, 4, 6, 8) # got too many positioned arguments
class ProxyDoesntUseKwOnly(object):
+
@frame_base.args_to_kwargs(Base)
@frame_base.populate_defaults(Base)
def func(self, a, **kwargs):
@@ -175,11 +187,14 @@ class FrameBaseTest(unittest.TestCase):
})
def test_populate_defaults_overwrites_copy(self):
+
class Base(object):
+
def func(self, a=1, b=2, c=3, *, copy=None):
pass
class Proxy(object):
+
@frame_base.args_to_kwargs(Base)
@frame_base.populate_defaults(Base)
def func(self, a, copy, **kwargs):
diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py
index ccd01f35f87..b2d22a58c06 100644
--- a/sdks/python/apache_beam/dataframe/frames.py
+++ b/sdks/python/apache_beam/dataframe/frames.py
@@ -65,6 +65,7 @@ PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:2]))
def populate_not_implemented(pd_type):
+
def wrapper(deferred_type):
for attr in dir(pd_type):
# Don't auto-define hidden methods or dunders
@@ -91,6 +92,7 @@ def populate_not_implemented(pd_type):
def _fillna_alias(method):
+
def wrapper(self, *args, **kwargs):
return self.fillna(*args, method=method, **kwargs)
@@ -153,6 +155,7 @@ UNLIFTABLE_GLOBAL_AGGREGATIONS = (
def _agg_method(base, func):
+
def wrapper(self, *args, **kwargs):
return self.agg(func, *args, **kwargs)
@@ -177,6 +180,7 @@ _PEEK_METHOD_EXPLANATION = (
class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
+
def _render_indexes(self):
if self.index.nlevels == 1:
return 'index=' + (
@@ -232,13 +236,9 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'drop',
- lambda df: df.drop(
- axis=axis,
- index=index,
- columns=columns,
- errors=errors,
- **kwargs), [self._expr],
+ 'drop', lambda df: df.drop(
+ axis=axis, index=index, columns=columns, errors=errors, **kwargs
+ ), [self._expr],
proxy=proxy,
requires_partition_by=requires))
@@ -248,8 +248,8 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
def droplevel(self, level, axis):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'droplevel',
- lambda df: df.droplevel(level, axis=axis), [self._expr],
+ 'droplevel', lambda df: df.droplevel(level, axis=axis),
+ [self._expr],
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Arbitrary()
if axis in (1, 'column') else partitionings.Singleton()))
@@ -259,8 +259,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
def swaplevel(self, **kwargs):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'swaplevel',
- lambda df: df.swaplevel(**kwargs), [self._expr],
+ 'swaplevel', lambda df: df.swaplevel(**kwargs), [self._expr],
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Arbitrary()))
@@ -295,13 +294,13 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
# This is OK, as its index must be the same size as the columns set of
# self, so cannot be too large.
class AsScalar(object):
+
def __init__(self, value):
self.value = value
with expressions.allow_non_parallel_operations():
value_expr = expressions.ComputedExpression(
- 'as_scalar',
- lambda df: AsScalar(df), [value._expr],
+ 'as_scalar', lambda df: AsScalar(df), [value._expr],
requires_partition_by=partitionings.Singleton())
get_value = lambda x: x.value
@@ -322,14 +321,9 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
return frame_base.DeferredFrame.wrap(
# yapf: disable
expressions.ComputedExpression(
- 'fillna',
- lambda df,
- value: df.fillna(
- get_value(value),
- method=method,
- axis=axis,
- limit=limit,
- **kwargs), [self._expr, value_expr],
+ 'fillna', lambda df, value: df.fillna(
+ get_value(value), method=method, axis=axis, limit=limit, **
+ kwargs), [self._expr, value_expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=requires))
@@ -345,30 +339,30 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
@frame_base.with_docs_from(pd.DataFrame)
def first(self, offset):
per_partition = expressions.ComputedExpression(
- 'first-per-partition',
- lambda df: df.sort_index().first(offset=offset), [self._expr],
+ 'first-per-partition', lambda df: df.sort_index().first(offset=offset),
+ [self._expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Arbitrary())
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'first',
- lambda df: df.sort_index().first(offset=offset), [per_partition],
+ 'first', lambda df: df.sort_index().first(offset=offset),
+ [per_partition],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Singleton()))
@frame_base.with_docs_from(pd.DataFrame)
def last(self, offset):
per_partition = expressions.ComputedExpression(
- 'last-per-partition',
- lambda df: df.sort_index().last(offset=offset), [self._expr],
+ 'last-per-partition', lambda df: df.sort_index().last(offset=offset),
+ [self._expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Arbitrary())
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'last',
- lambda df: df.sort_index().last(offset=offset), [per_partition],
+ 'last', lambda df: df.sort_index().last(offset=offset),
+ [per_partition],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Singleton()))
@@ -388,8 +382,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
if axis in (1, 'columns'):
return _DeferredGroupByCols(
expressions.ComputedExpression(
- 'groupbycols',
- lambda df: df.groupby(
+ 'groupbycols', lambda df: df.groupby(
by, axis=axis, group_keys=group_keys, **kwargs), [self._expr],
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Arbitrary()),
@@ -561,11 +554,9 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
return DeferredGroupBy(
expressions.ComputedExpression(
- 'groupbyindex',
- lambda df: df.groupby(
- level=list(range(df.index.nlevels)),
- group_keys=group_keys,
- **kwargs), [to_group],
+ 'groupbyindex', lambda df: df.groupby(
+ level=list(range(df.index.nlevels)), group_keys=group_keys, **
+ kwargs), [to_group],
requires_partition_by=partitionings.Index(),
preserves_partition_by=partitionings.Arbitrary()),
kwargs,
@@ -612,8 +603,8 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
requires_partition_by = partitionings.Arbitrary()
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'reset_index',
- lambda df: df.reset_index(level=level, **kwargs), [self._expr],
+ 'reset_index', lambda df: df.reset_index(level=level, **kwargs),
+ [self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=requires_partition_by))
@@ -705,12 +696,8 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
"requires collecting all data on a single node."))
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'replace',
- lambda df: df.replace(
- to_replace=to_replace,
- value=value,
- limit=limit,
- method=method,
+ 'replace', lambda df: df.replace(
+ to_replace=to_replace, value=value, limit=limit, method=method,
**kwargs), [self._expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=requires_partition_by))
@@ -732,10 +719,8 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
elif isinstance(ambiguous, frame_base.DeferredFrame):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'tz_localize',
- lambda df,
- ambiguous: df.tz_localize(ambiguous=ambiguous, **kwargs),
- [self._expr, ambiguous._expr],
+ 'tz_localize', lambda df, ambiguous: df.tz_localize(
+ ambiguous=ambiguous, **kwargs), [self._expr, ambiguous._expr],
requires_partition_by=partitionings.Index(),
preserves_partition_by=partitionings.Singleton()))
elif ambiguous == 'infer':
@@ -767,8 +752,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'sum_sizes',
- lambda sizes: sizes.sum(), [sizes],
+ 'sum_sizes', lambda sizes: sizes.sum(), [sizes],
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -787,8 +771,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'sum_lengths',
- lambda lengths: lengths.sum(), [lengths],
+ 'sum_lengths', lambda lengths: lengths.sum(), [lengths],
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -812,8 +795,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'check_all_empty',
- lambda empties: empties.all(), [empties],
+ 'check_all_empty', lambda empties: empties.all(), [empties],
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -834,8 +816,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
# Will throw if overall dataset has != 1 element
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'combine_all_bools',
- lambda bools: bools.bool(), [bools],
+ 'combine_all_bools', lambda bools: bools.bool(), [bools],
proxy=bool(),
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -845,8 +826,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
intermediate = expressions.ComputedExpression(
'equals_partitioned',
# Wrap scalar results in a Series for easier concatenation later
- lambda df,
- other: pd.Series(df.equals(other)),
+ lambda df, other: pd.Series(df.equals(other)),
[self._expr, other._expr],
requires_partition_by=partitionings.Index(),
preserves_partition_by=partitionings.Singleton())
@@ -854,8 +834,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'aggregate_equals',
- lambda df: df.all(), [intermediate],
+ 'aggregate_equals', lambda df: df.all(), [intermediate],
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -1010,8 +989,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
"Please upgrade to pandas 1.2.0 or higher to use this operation.")
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'unstack',
- lambda s: s.unstack(**kwargs), [self._expr],
+ 'unstack', lambda s: s.unstack(**kwargs), [self._expr],
requires_partition_by=partitionings.Index()))
else:
# Unstacking MultiIndex objects
@@ -1062,8 +1040,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'unstack',
- lambda s: pd.concat([proxy, s.unstack(**kwargs)]),
+ 'unstack', lambda s: pd.concat([proxy, s.unstack(**kwargs)]),
[self._expr],
proxy=proxy,
requires_partition_by=partitionings.Singleton()))
@@ -1080,8 +1057,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
# KeyError at construction time for missing columns.
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'xs',
- lambda df: df.xs(key, axis=axis, **kwargs), [self._expr],
+ 'xs', lambda df: df.xs(key, axis=axis, **kwargs), [self._expr],
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Arbitrary()))
elif axis not in ('index', 0):
@@ -1239,6 +1215,7 @@ class DeferredDataFrameOrSeries(frame_base.DeferredFrame):
@populate_not_implemented(pd.Series)
@frame_base.DeferredFrame._register_for(pd.Series)
class DeferredSeries(DeferredDataFrameOrSeries):
+
def __repr__(self):
return (
f'DeferredSeries(name={self.name!r}, dtype={self.dtype}, '
@@ -1251,6 +1228,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
@name.setter
def name(self, value):
+
def fn(s):
s = s.copy()
s.name = value
@@ -1266,16 +1244,14 @@ class DeferredSeries(DeferredDataFrameOrSeries):
@frame_base.with_docs_from(pd.Series)
def hasnans(self):
has_nans = expressions.ComputedExpression(
- 'hasnans',
- lambda s: pd.Series(s.hasnans), [self._expr],
+ 'hasnans', lambda s: pd.Series(s.hasnans), [self._expr],
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Singleton())
with expressions.allow_non_parallel_operations():
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'combine_hasnans',
- lambda s: s.any(), [has_nans],
+ 'combine_hasnans', lambda s: s.any(), [has_nans],
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -1312,8 +1288,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
expressions.ComputedExpression(
# yapf: disable
'getitem',
- lambda df,
- indexer: df[indexer],
+ lambda df, indexer: df[indexer],
[self._expr, key._expr],
requires_partition_by=partitionings.Index(),
preserves_partition_by=partitionings.Arbitrary()))
@@ -1369,9 +1344,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'append',
- lambda s,
- to_append: s.append(
+ 'append', lambda s, to_append: s.append(
to_append, verify_integrity=verify_integrity, **kwargs),
[self._expr, to_append._expr],
requires_partition_by=requires,
@@ -1399,9 +1372,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
# multiple return values.
aligned = frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'align',
- lambda x,
- y: pd.concat([x, y], axis=1, join='inner'),
+ 'align', lambda x, y: pd.concat([x, y], axis=1, join='inner'),
[self._expr, other._expr],
requires_partition_by=partitionings.Index(),
preserves_partition_by=partitionings.Arbitrary()))
@@ -1488,8 +1459,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'idx_combine',
- lambda s: func(s, **kwargs), [idx_func],
+ 'idx_combine', lambda s: func(s, **kwargs), [idx_func],
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -1514,8 +1484,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
partitionings.Singleton() if ignore_index else partitionings.Index())
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'explode',
- lambda s: s.explode(ignore_index), [self._expr],
+ 'explode', lambda s: s.explode(ignore_index), [self._expr],
preserves_partition_by=preserves,
requires_partition_by=partitionings.Arbitrary()))
@@ -1558,8 +1527,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
if right_is_series:
result = expressions.ComputedExpression(
- 'extract',
- lambda df: df[0], [sums],
+ 'extract', lambda df: df[0], [sums],
requires_partition_by=partitionings.Singleton())
else:
result = sums
@@ -1592,8 +1560,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'quantile',
- lambda df: df.quantile(q=q, **kwargs), [self._expr],
+ 'quantile', lambda df: df.quantile(q=q, **kwargs), [self._expr],
requires_partition_by=requires,
preserves_partition_by=partitionings.Singleton()))
@@ -1679,9 +1646,8 @@ class DeferredSeries(DeferredDataFrameOrSeries):
# and custom partitioning.
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'corr',
- lambda df,
- other: df.corr(other, method=method, min_periods=min_periods),
+ 'corr', lambda df, other: df.corr(
+ other, method=method, min_periods=min_periods),
[self._expr, other._expr],
requires_partition_by=partitionings.Singleton(reason=reason)))
@@ -1896,8 +1862,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
def dropna(self, **kwargs):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'dropna',
- lambda df: df.dropna(**kwargs), [self._expr],
+ 'dropna', lambda df: df.dropna(**kwargs), [self._expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Arbitrary()))
@@ -2006,8 +1971,8 @@ class DeferredSeries(DeferredDataFrameOrSeries):
rows = [self.agg([f], *args, **kwargs) for f in func]
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'join_aggregate',
- lambda *rows: pd.concat(rows), [row._expr for row in rows]))
+ 'join_aggregate', lambda *rows: pd.concat(rows),
+ [row._expr for row in rows]))
else:
# We're only handling a single column. It could be 'func' or ['func'],
# which produce different results. 'func' produces a scalar, ['func']
@@ -2190,15 +2155,14 @@ class DeferredSeries(DeferredDataFrameOrSeries):
reason="order-sensitive")
kwargs['keep'] = keep
per_partition = expressions.ComputedExpression(
- 'nlargest-per-partition',
- lambda df: df.nlargest(**kwargs), [self._expr],
+ 'nlargest-per-partition', lambda df: df.nlargest(**kwargs),
+ [self._expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Arbitrary())
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'nlargest',
- lambda df: df.nlargest(**kwargs), [per_partition],
+ 'nlargest', lambda df: df.nlargest(**kwargs), [per_partition],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Singleton()))
@@ -2220,21 +2184,21 @@ class DeferredSeries(DeferredDataFrameOrSeries):
reason="order-sensitive")
kwargs['keep'] = keep
per_partition = expressions.ComputedExpression(
- 'nsmallest-per-partition',
- lambda df: df.nsmallest(**kwargs), [self._expr],
+ 'nsmallest-per-partition', lambda df: df.nsmallest(**kwargs),
+ [self._expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Arbitrary())
with expressions.allow_non_parallel_operations(True):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'nsmallest',
- lambda df: df.nsmallest(**kwargs), [per_partition],
+ 'nsmallest', lambda df: df.nsmallest(**kwargs), [per_partition],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Singleton()))
@property # type: ignore
@frame_base.with_docs_from(pd.Series)
def is_unique(self):
+
def set_index(s):
s = s[:]
s.index = s
@@ -2247,16 +2211,14 @@ class DeferredSeries(DeferredDataFrameOrSeries):
preserves_partition_by=partitionings.Singleton())
is_unique_distributed = expressions.ComputedExpression(
- 'is_unique_distributed',
- lambda s: pd.Series(s.is_unique), [self_index],
+ 'is_unique_distributed', lambda s: pd.Series(s.is_unique), [self_index],
requires_partition_by=partitionings.Index(),
preserves_partition_by=partitionings.Singleton())
with expressions.allow_non_parallel_operations():
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'combine',
- lambda s: s.all(), [is_unique_distributed],
+ 'combine', lambda s: s.all(), [is_unique_distributed],
requires_partition_by=partitionings.Singleton(),
preserves_partition_by=partitionings.Singleton()))
@@ -2293,8 +2255,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
reason="non-deferred-result")
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'unique',
- lambda df: pd.Series(df.unique()), [self._expr],
+ 'unique', lambda df: pd.Series(df.unique()), [self._expr],
preserves_partition_by=partitionings.Singleton(),
requires_partition_by=partitionings.Singleton(
reason="unique() cannot currently be parallelized.")))
@@ -2302,9 +2263,8 @@ class DeferredSeries(DeferredDataFrameOrSeries):
@frame_base.with_docs_from(pd.Series)
def update(self, other):
self._expr = expressions.ComputedExpression(
- 'update',
- lambda df,
- other: df.update(other) or df, [self._expr, other._expr],
+ 'update', lambda df, other: df.update(other) or df,
+ [self._expr, other._expr],
preserves_partition_by=partitionings.Arbitrary(),
requires_partition_by=partitionings.Index())
@@ -2424,16 +2384,14 @@ class DeferredSeries(DeferredDataFrameOrSeries):
if isinstance(repeats, int):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
- 'repeat',
- lambda series: series.repeat(repeats), [self._expr],
+ 'repeat', lambda series: series.repeat(repeats), [self._expr],
requires_partition_by=partitionings.Arbitrary(),
preserves_partition_by=partitionings.Arbitrary()))
elif isinstance(repeats, frame_base.DeferredBase):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'repeat',
- lambda series,
- repeats_series: series.repeat(repeats_series),
+ lambda series, repeats_series: series.repeat(repeats_series),
[self._expr, repeats._expr],
requires_partition_by=partitionings.Index(),
preserves_partition_by=partitionings.Arbitrary()))
@@ -2467,8 +2425,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
return frame_base.DeferredFrame.wrap(
expressions.ComputedExpression(
'compare',
- lambda s,
- other: s.compare(other, align_axis, **kwargs),
+ lambda s, other: s.compare(other, align_axis, **kwargs),
[self._expr, other._expr],
requires_partition_by=partitionings.Index(),
preserves_partition_by=preserves_partition))
@@ -2477,6 +2434,7 @@ class DeferredSeries(DeferredDataFrameOrSeries):
@populate_not_implemented(pd.DataFrame)
@frame_base.DeferredFrame._register_for(pd.DataFrame)
class DeferredDataFrame(DeferredDataFrameOrSeries):
+
def __repr__(self):
return (
f'DeferredDataFrame(columns={list(self.columns)}, '
@@ -2489,6 +2447,7 @@ class DeferredDataFrame(DeferredDataFrameOrSeries):
@columns.setter
def columns(self, columns):
+
def set_columns(df):
df = df.copy()
df.columns = columns
diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py
index f99b77e446a..d739152f6fd 100644
--- a/sdks/python/apache_beam/dataframe/frames_test.py
+++ b/sdks/python/apache_beam/dataframe/frames_test.py
@@ -61,6 +61,7 @@ def _get_deferred_args(*args):
class _AbstractFrameTest(unittest.TestCase):
"""Test sub-class with utilities for verifying DataFrame operations."""
+
def _run_error_test(
self, func, *args, construction_time=True, distributed=True):
"""Verify that func(*args) raises the same exception in pandas and in Beam.
@@ -119,6 +120,7 @@ class _AbstractFrameTest(unittest.TestCase):
Checks that func performs the same inplace operation on arg, in pandas and
in Beam."""
+
def wrapper(df):
df = df.copy()
func(df)
@@ -273,6 +275,7 @@ class _AbstractFrameTest(unittest.TestCase):
class DeferredFrameTest(_AbstractFrameTest):
"""Miscellaneous tessts for DataFrame operations."""
+
def test_series_arithmetic(self):
a = pd.Series([1, 2, 3])
b = pd.Series([100, 200, 300])
@@ -345,6 +348,7 @@ class DeferredFrameTest(_AbstractFrameTest):
lambda df: df.xs('state'), df.set_index(['provider', 'time']))
def test_set_column(self):
+
def new_column(df):
df['NewCol'] = df['Speed']
@@ -355,6 +359,7 @@ class DeferredFrameTest(_AbstractFrameTest):
self._run_inplace_test(new_column, df)
def test_set_column_from_index(self):
+
def new_column(df):
df['NewCol'] = df.index
@@ -390,8 +395,7 @@ class DeferredFrameTest(_AbstractFrameTest):
ambiguous = pd.Series([True, True, False], index=s.index)
self._run_test(
- lambda s,
- ambiguous: s.tz_localize('CET', ambiguous=ambiguous),
+ lambda s, ambiguous: s.tz_localize('CET', ambiguous=ambiguous),
s,
ambiguous)
@@ -444,8 +448,7 @@ class DeferredFrameTest(_AbstractFrameTest):
df2 = pd.DataFrame({'A': [1, 1], 'B': [3, 3]})
take_smaller = lambda s1, s2: s1 if s1.sum() < s2.sum() else s2
self._run_test(
- lambda df,
- df2: df.combine(df2, take_smaller),
+ lambda df, df2: df.combine(df2, take_smaller),
df,
df2,
nonparallel=True)
@@ -455,8 +458,7 @@ class DeferredFrameTest(_AbstractFrameTest):
df2 = pd.DataFrame({'A': [1, 1], 'B': [3, 3]})
take_smaller = lambda s1, s2: s1 if s1.sum() < s2.sum() else s2
self._run_test(
- lambda df1,
- df2: df1.combine(df2, take_smaller, fill_value=-5),
+ lambda df1, df2: df1.combine(df2, take_smaller, fill_value=-5),
df1,
df2,
nonparallel=True)
@@ -465,8 +467,7 @@ class DeferredFrameTest(_AbstractFrameTest):
s1 = pd.Series({'falcon': 330.0, 'eagle': 160.0})
s2 = pd.Series({'falcon': 345.0, 'eagle': 200.0, 'duck': 30.0})
self._run_test(
- lambda s1,
- s2: s1.combine(s2, max),
+ lambda s1, s2: s1.combine(s2, max),
s1,
s2,
nonparallel=True,
@@ -576,16 +577,14 @@ class DeferredFrameTest(_AbstractFrameTest):
'rkey': ['foo', 'bar', 'baz', 'foo'], 'value': [5, 6, 7, 8]
})
self._run_test(
- lambda df1,
- df2: df1.merge(df2, left_on='lkey', right_on='rkey').rename(
+ lambda df1, df2: df1.merge(df2, left_on='lkey', right_on='rkey').rename(
index=lambda x: '*'),
df1,
df2,
nonparallel=True,
check_proxy=False)
self._run_test(
- lambda df1,
- df2: df1.merge(
+ lambda df1, df2: df1.merge(
df2, left_on='lkey', right_on='rkey', suffixes=('_left', '_right')).
rename(index=lambda x: '*'),
df1,
@@ -600,8 +599,8 @@ class DeferredFrameTest(_AbstractFrameTest):
df2 = pd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]})
self._run_test(
- lambda df1,
- df2: df1.merge(df2, how='left', on='a').rename(index=lambda x: '*'),
+ lambda df1, df2: df1.merge(df2, how='left', on='a').rename(
+ index=lambda x: '*'),
df1,
df2,
nonparallel=True,
@@ -618,8 +617,7 @@ class DeferredFrameTest(_AbstractFrameTest):
}).set_index('rkey')
self._run_test(
- lambda df1,
- df2: df1.merge(df2, left_index=True, right_index=True),
+ lambda df1, df2: df1.merge(df2, left_index=True, right_index=True),
df1,
df2,
check_proxy=False)
@@ -632,16 +630,14 @@ class DeferredFrameTest(_AbstractFrameTest):
'key': ['foo', 'bar', 'baz', 'foo'], 'value': [5, 6, 7, 8]
})
self._run_test(
- lambda df1,
- df2: df1.merge(df2, on='key').rename(index=lambda x: '*'),
+ lambda df1, df2: df1.merge(df2, on='key').rename(index=lambda x: '*'),
df1,
df2,
nonparallel=True,
check_proxy=False)
self._run_test(
- lambda df1,
- df2: df1.merge(df2, on='key', suffixes=('_left', '_right')).rename(
- index=lambda x: '*'),
+ lambda df1, df2: df1.merge(df2, on='key', suffixes=('_left', '_right')).
+ rename(index=lambda x: '*'),
df1,
df2,
nonparallel=True,
@@ -652,16 +648,15 @@ class DeferredFrameTest(_AbstractFrameTest):
df2 = pd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]})
self._run_test(
- lambda df1,
- df2: df1.merge(df2, how='left', on='a').rename(index=lambda x: '*'),
+ lambda df1, df2: df1.merge(df2, how='left', on='a').rename(
+ index=lambda x: '*'),
df1,
df2,
nonparallel=True,
check_proxy=False)
# Test without specifying 'on'
self._run_test(
- lambda df1,
- df2: df1.merge(df2, how='left').rename(index=lambda x: '*'),
+ lambda df1, df2: df1.merge(df2, how='left').rename(index=lambda x: '*'),
df1,
df2,
nonparallel=True,
@@ -672,8 +667,7 @@ class DeferredFrameTest(_AbstractFrameTest):
df2 = pd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4], 'a_rsuffix': [7, 8]})
self._run_test(
- lambda df1,
- df2: df1.merge(
+ lambda df1, df2: df1.merge(
df2, how='left', on='a', suffixes=('_lsuffix', '_rsuffix')).rename(
index=lambda x: '*'),
df1,
@@ -682,9 +676,9 @@ class DeferredFrameTest(_AbstractFrameTest):
check_proxy=False)
# Test without specifying 'on'
self._run_test(
- lambda df1,
- df2: df1.merge(df2, how='left', suffixes=('_lsuffix', '_rsuffix')).
- rename(index=lambda x: '*'),
+ lambda df1, df2: df1.merge(
+ df2, how='left', suffixes=('_lsuffix', '_rsuffix')).rename(
+ index=lambda x: '*'),
df1,
df2,
nonparallel=True,
@@ -731,10 +725,8 @@ class DeferredFrameTest(_AbstractFrameTest):
for normalize in (True, False):
for dropna in (True, False):
self._run_test(
- lambda df,
- dropna=dropna,
- normalize=normalize: df.num_wings.value_counts(
- dropna=dropna, normalize=normalize),
+ lambda df, dropna=dropna, normalize=normalize: df.num_wings.
+ value_counts(dropna=dropna, normalize=normalize),
df)
def test_value_counts_does_not_support_sort(self):
@@ -962,11 +954,8 @@ class DeferredFrameTest(_AbstractFrameTest):
df)
self._run_test(
lambda df: df.melt(
- id_vars=['A'],
- value_vars=['B'],
- var_name='myVarname',
- value_name='myValname',
- ignore_index=False),
+ id_vars=['A'], value_vars=['B'], var_name='myVarname', value_name=
+ 'myValname', ignore_index=False),
df)
self._run_test(
lambda df: df.melt(
@@ -1041,14 +1030,12 @@ class DeferredFrameTest(_AbstractFrameTest):
df2 = pd.DataFrame({'A': range(10), 'B': range(10)}, index=range(9, 19))
self._run_error_test(
- lambda s1,
- s2: s1.append(s2, verify_integrity=True),
+ lambda s1, s2: s1.append(s2, verify_integrity=True),
df1['A'],
df2['A'],
construction_time=False)
self._run_error_test(
- lambda df1,
- df2: df1.append(df2, verify_integrity=True),
+ lambda df1, df2: df1.append(df2, verify_integrity=True),
df1,
df2,
construction_time=False)
@@ -1144,12 +1131,12 @@ class DeferredFrameTest(_AbstractFrameTest):
(
lambda base: base.from_dict({
'row_1': [3, 2, 1, 0], 'row_2': ['a', 'b', 'c', 'd']
- },
- orient='index'), ),
+ }, orient='index'), ),
(
lambda base: base.from_records(
- np.array([(3, 'a'), (2, 'b'), (1, 'c'), (0, 'd')],
- dtype=[('col_1', 'i4'), ('col_2', 'U1')])), ),
+ np.array([(3, 'a'), (2, 'b'), (1, 'c'),
+ (0, 'd')], dtype=[('col_1', 'i4'),
+ ('col_2', 'U1')])), ),
])
def test_create_methods(self, func):
expected = func(pd.DataFrame)
@@ -1242,8 +1229,7 @@ class DeferredFrameTest(_AbstractFrameTest):
ambiguous = pd.Series([True, True, False], index=s.index)
self._run_test(
- lambda s,
- ambiguous: s.dt.tz_localize('CET', ambiguous=ambiguous),
+ lambda s, ambiguous: s.dt.tz_localize('CET', ambiguous=ambiguous),
s,
ambiguous)
@@ -1297,25 +1283,22 @@ class DeferredFrameTest(_AbstractFrameTest):
self._run_test(lambda df1, df2: df1.compare(df2), df1, df2)
self._run_test(
- lambda df1,
- df2: df1.compare(df2, align_axis=0),
+ lambda df1, df2: df1.compare(df2, align_axis=0),
df1,
df2,
check_proxy=False)
self._run_test(lambda df1, df2: df1.compare(df2, keep_shape=True), df1, df2)
self._run_test(
- lambda df1,
- df2: df1.compare(df2, align_axis=0, keep_shape=True),
+ lambda df1, df2: df1.compare(df2, align_axis=0, keep_shape=True),
df1,
df2)
self._run_test(
- lambda df1,
- df2: df1.compare(df2, keep_shape=True, keep_equal=True),
+ lambda df1, df2: df1.compare(df2, keep_shape=True, keep_equal=True),
df1,
df2)
self._run_test(
- lambda df1,
- df2: df1.compare(df2, align_axis=0, keep_shape=True, keep_equal=True),
+ lambda df1, df2: df1.compare(
+ df2, align_axis=0, keep_shape=True, keep_equal=True),
df1,
df2)
@@ -1398,6 +1381,7 @@ class DeferredFrameTest(_AbstractFrameTest):
self._run_test(lambda s2: s2.idxmax(skipna=False), s2)
def test_pipe(self):
+
def df_times(df, column, times):
df[column] = df[column] * times
return df
@@ -1720,6 +1704,7 @@ def numeric_only_kwargs_for_pandas_2(agg_type: str) -> dict[str, bool]:
class GroupByTest(_AbstractFrameTest):
"""Tests for DataFrame/Series GroupBy operations."""
+
@staticmethod
def median_sum_fn(x):
with warnings.catch_warnings():
@@ -1956,8 +1941,7 @@ class GroupByTest(_AbstractFrameTest):
[2, 1],
['foo', 0],
[1, 'str'],
- [3, 0, 2, 1],
- ])
+ [3, 0, 2, 1], ])
def test_groupby_level_agg(self, level):
df = GROUPBY_DF.set_index(['group', 'foo', 'bar', 'str'], drop=False)
self._run_test(lambda df: df.groupby(level=level).bar.max(), df)
@@ -2212,8 +2196,8 @@ class AggregationTest(_AbstractFrameTest):
def test_series_agg_multifunc_level(self):
# level= is ignored for multiple agg fns
self._run_test(
- lambda df: df.set_index(['group', 'foo']).bar.agg(['min', 'max'],
- level=0),
+ lambda df: df.set_index(['group', 'foo']).bar.agg(['min', 'max'], level=
+ 0),
GROUPBY_DF)
def test_series_mean_skipna(self):
@@ -2408,6 +2392,7 @@ class BeamSpecificTest(unittest.TestCase):
"""Tests for functionality that's specific to the Beam DataFrame API.
These features don't exist in pandas so we must verify them independently."""
+
def assert_frame_data_equivalent(
self, actual, expected, check_column_subset=False, extra_col_value=0):
"""Verify that actual is the same as expected, ignoring the index and order
@@ -2807,8 +2792,7 @@ class BeamSpecificTest(unittest.TestCase):
self.assertTrue(target_weight > other_weight * 10, "weights too close")
result = self._evaluate(
- lambda s,
- weights: s.sample(n=num_samples, weights=weights).sum(),
+ lambda s, weights: s.sample(n=num_samples, weights=weights).sum(),
# The first elements are 1, the rest are all 0. This means that when
# we sum all the sampled elements (above), the result should be the
# number of times the first elements (aka targets) were sampled.
@@ -2941,6 +2925,7 @@ class BeamSpecificTest(unittest.TestCase):
class AllowNonParallelTest(unittest.TestCase):
+
def _use_non_parallel_operation(self):
_ = frame_base.DeferredFrame.wrap(
expressions.PlaceholderExpression(pd.Series([1, 2, 3]))).replace(
@@ -3036,6 +3021,7 @@ class ConstructionTimeTest(unittest.TestCase):
class DocstringTest(unittest.TestCase):
+
@parameterized.expand([
(frames.DeferredDataFrame, pd.DataFrame),
(frames.DeferredSeries, pd.Series),
@@ -3073,6 +3059,7 @@ class DocstringTest(unittest.TestCase):
class ReprTest(unittest.TestCase):
+
def test_basic_dataframe(self):
df = frame_base.DeferredFrame.wrap(
expressions.ConstantExpression(GROUPBY_DF))
@@ -3188,6 +3175,7 @@ class ReprTest(unittest.TestCase):
'[interactive] dependency is not installed.')
@isolated_env
class InteractiveDataFrameTest(unittest.TestCase):
+
def test_collect_merged_dataframes(self):
p = beam.Pipeline(InteractiveRunner())
pcoll_1 = (
diff --git a/sdks/python/apache_beam/dataframe/io.py b/sdks/python/apache_beam/dataframe/io.py
index 5fcb7326a02..d7a5e650b59 100644
--- a/sdks/python/apache_beam/dataframe/io.py
+++ b/sdks/python/apache_beam/dataframe/io.py
@@ -171,8 +171,7 @@ def to_json(df, path, orient=None, *args, **kwargs):
@frame_base.with_docs_from(pd)
def read_html(path, *args, **kwargs):
return _ReadFromPandas(
- lambda *args,
- **kwargs: pd.read_html(*args, **kwargs)[0],
+ lambda *args, **kwargs: pd.read_html(*args, **kwargs)[0],
path,
args,
kwargs)
@@ -193,8 +192,8 @@ def to_html(df, path, *args, **kwargs):
def _binary_reader(format):
func = getattr(pd, 'read_%s' % format)
- result = lambda path, *args, **kwargs: _ReadFromPandas(func, path, args,
- kwargs)
+ result = lambda path, *args, **kwargs: _ReadFromPandas(
+ func, path, args, kwargs)
result.__name__ = f'read_{format}'
return result
@@ -202,10 +201,8 @@ def _binary_reader(format):
def _binary_writer(format):
result = (
- lambda df,
- path,
- *args,
- **kwargs: _as_pc(df) | _WriteToPandas(f'to_{format}', path, args, kwargs))
+ lambda df, path, *args, **kwargs: _as_pc(df) | _WriteToPandas(
+ f'to_{format}', path, args, kwargs))
result.__name__ = f'to_{format}'
return result
@@ -249,6 +246,7 @@ def _shift_range_index(offset, df):
class _ReadFromPandas(beam.PTransform):
+
def __init__(
self,
reader,
@@ -294,9 +292,10 @@ class _ReadFromPandas(beam.PTransform):
matches_pcoll.pipeline
| 'DoOnce' >> beam.Create([None])
| beam.Map(
- lambda _,
- paths: {path: ix
- for ix, path in enumerate(sorted(paths))},
+ lambda _, paths: {
+ path: ix
+ for ix, path in enumerate(sorted(paths))
+ },
paths=beam.pvalue.AsList(
matches_pcoll | beam.Map(lambda match: match.path))))
@@ -318,6 +317,7 @@ class _ReadFromPandas(beam.PTransform):
class _Splitter:
+
def empty_buffer(self):
"""Returns an empty buffer of the right type (string or bytes).
"""
@@ -351,6 +351,7 @@ class _DelimSplitter(_Splitter):
This delimiter is assumed ot never occur within a record.
"""
+
def __init__(self, delim, read_chunk_size=_DEFAULT_BYTES_CHUNKSIZE):
# Multi-char delimiters would require more care across chunk boundaries.
assert len(delim) == 1
@@ -393,6 +394,7 @@ class _TextFileSplitter(_DelimSplitter):
Currently does not handle quoted newlines, so is off by default, but such
support could be added in the future.
"""
+
def __init__(self, args, kwargs, read_chunk_size=_DEFAULT_BYTES_CHUNKSIZE):
if args:
# TODO(robertwb): Automatically populate kwargs as we do for df methods.
@@ -470,6 +472,7 @@ class _TruncatingFileHandle(object):
As with all SDF trackers, the endpoint may change dynamically during reading.
"""
+
def __init__(self, underlying, tracker, splitter):
self._underlying = underlying
self._tracker = tracker
@@ -582,6 +585,7 @@ class _TruncatingFileHandle(object):
class _ReadFromPandasDoFn(beam.DoFn, beam.RestrictionProvider):
+
def __init__(self, reader, args, kwargs, binary, incremental, splitter):
# avoid pickling issues
if reader.__module__.startswith('pandas.'):
@@ -653,6 +657,7 @@ class _ReadFromPandasDoFn(beam.DoFn, beam.RestrictionProvider):
class _WriteToPandas(beam.PTransform):
+
def __init__(
self, writer, path, args, kwargs, incremental=False, binary=True):
self.writer = writer
@@ -677,6 +682,7 @@ class _WriteToPandas(beam.PTransform):
class _WriteToPandasFileSink(fileio.FileSink):
+
def __init__(self, writer, args, kwargs, incremental, binary):
if 'compression' in kwargs:
raise NotImplementedError('compression')
@@ -765,6 +771,7 @@ class _WriteToPandasFileSink(fileio.FileSink):
class ReadViaPandas(beam.PTransform):
+
def __init__(
self,
format,
@@ -787,6 +794,7 @@ class ReadViaPandas(beam.PTransform):
class WriteViaPandas(beam.PTransform):
+
def __init__(self, format, *args, **kwargs):
self._writer_func = globals()['to_%s' % format]
self._args = args
@@ -830,6 +838,7 @@ class _ReadGbq(beam.PTransform):
unspecified or set to false, the default is currently utilized (EXPORT).
If the flag is set to true,
'DIRECT_READ' will be utilized."""
+
def __init__(
self,
table=None,
diff --git a/sdks/python/apache_beam/dataframe/io_it_test.py b/sdks/python/apache_beam/dataframe/io_it_test.py
index 9f750e2ff58..3e4a9a037f9 100644
--- a/sdks/python/apache_beam/dataframe/io_it_test.py
+++ b/sdks/python/apache_beam/dataframe/io_it_test.py
@@ -40,6 +40,7 @@ except ImportError:
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class ReadUsingReadGbqTests(unittest.TestCase):
+
@pytest.mark.it_postcommit
def test_ReadGbq(self):
from apache_beam.dataframe import convert
diff --git a/sdks/python/apache_beam/dataframe/io_test.py b/sdks/python/apache_beam/dataframe/io_test.py
index 92bb10225c7..17784589c66 100644
--- a/sdks/python/apache_beam/dataframe/io_test.py
+++ b/sdks/python/apache_beam/dataframe/io_test.py
@@ -69,6 +69,7 @@ class MyRow(typing.NamedTuple):
platform.system() == 'Windows',
'https://github.com/apache/beam/issues/20642')
class IOTest(unittest.TestCase):
+
def setUp(self):
self._temp_roots = []
@@ -200,6 +201,7 @@ A B
big['text'] = big.number.map(lambda n: 'f' + 'o' * n)
def frame_equal_to(expected_, check_index=True, check_names=True):
+
def check(actual):
expected = expected_
try:
@@ -312,6 +314,7 @@ A B
('skiprows', dict(skiprows=[0, 1], header=[0, 1], comment='X')),
])
def test_csv_splitter(self, name, kwargs):
+
def assert_frame_equal(expected, actual):
try:
pandas.testing.assert_frame_equal(expected, actual)
@@ -442,6 +445,7 @@ X , c1, c2
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class ReadGbqTransformTests(unittest.TestCase):
+
@mock.patch.object(BigQueryWrapper, 'get_table')
def test_bad_schema_public_api_direct_read(self, get_table):
try:
@@ -466,6 +470,7 @@ class ReadGbqTransformTests(unittest.TestCase):
table="dataset.sample_table", use_bqstorage_api=True)
def test_unsupported_callable(self):
+
def filterTable(table):
if table is not None:
return table
diff --git a/sdks/python/apache_beam/dataframe/pandas_doctests_test.py b/sdks/python/apache_beam/dataframe/pandas_doctests_test.py
index c7ea908a933..70e8693ab17 100644
--- a/sdks/python/apache_beam/dataframe/pandas_doctests_test.py
+++ b/sdks/python/apache_beam/dataframe/pandas_doctests_test.py
@@ -27,6 +27,7 @@ from apache_beam.dataframe.pandas_top_level_functions import _is_top_level_funct
@unittest.skipIf(
sys.platform == 'win32', '[https://github.com/apache/beam/issues/20361]')
class DoctestTest(unittest.TestCase):
+
def test_ndframe_tests(self):
# IO methods are tested in io_test.py
skip_writes = {
@@ -70,8 +71,7 @@ class DoctestTest(unittest.TestCase):
"df.loc['2016-01-05':'2016-01-10', :].tail()"
],
'pandas.core.generic.NDFrame.replace': [
- "s.replace([1, 2], method='bfill')",
- # Relies on method='pad'
+ "s.replace([1, 2], method='bfill')", # Relies on method='pad'
"s.replace('a')",
# Relies on method='pad'
# value=None is not valid for pandas < 1.4
@@ -142,8 +142,7 @@ class DoctestTest(unittest.TestCase):
# some kind, 2 was passed
# pandas doctests only verify the type of exception
'df.rename(2)'
- ],
- # For pandas >= 1.4, rename is changed to _rename
+ ], # For pandas >= 1.4, rename is changed to _rename
'pandas.core.generic.NDFrame._rename': [
# Seems to be an upstream bug. The actual error has a different
# message:
@@ -151,8 +150,7 @@ class DoctestTest(unittest.TestCase):
# some kind, 2 was passed
# pandas doctests only verify the type of exception
'df.rename(2)'
- ],
- # Tests rely on setting index
+ ], # Tests rely on setting index
'pandas.core.generic.NDFrame.rename_axis': ['*'],
# Raises right exception, but testing framework has matching issues.
'pandas.core.generic.NDFrame.replace': [
@@ -160,7 +158,7 @@ class DoctestTest(unittest.TestCase):
],
'pandas.core.generic.NDFrame.squeeze': ['*'],
- # NameError
+ # NameError
'pandas.core.generic.NDFrame.resample': ['df'],
# Skipped so we don't need to install natsort
@@ -211,8 +209,7 @@ class DoctestTest(unittest.TestCase):
"df.nsmallest(3, 'population', keep='last')",
],
'pandas.core.frame.DataFrame.replace': [
- "s.replace([1, 2], method='bfill')",
- # Relies on method='pad'
+ "s.replace([1, 2], method='bfill')", # Relies on method='pad'
"s.replace('a')",
# Relies on method='pad'
# value=None is not valid for pandas < 1.4
@@ -256,8 +253,7 @@ class DoctestTest(unittest.TestCase):
"df.melt(id_vars=[('A', 'D')], value_vars=[('B', 'E')])",
"df.melt(id_vars=['A'], value_vars=['B'],\n" +
" var_name='myVarname', value_name='myValname')"
- ],
- # Most keep= options are order-sensitive
+ ], # Most keep= options are order-sensitive
'pandas.core.frame.DataFrame.drop_duplicates': ['*'],
'pandas.core.frame.DataFrame.duplicated': [
'df.duplicated()',
@@ -294,20 +290,18 @@ class DoctestTest(unittest.TestCase):
"df1.merge(df2, how='cross')"
],
- # TODO(https://github.com/apache/beam/issues/20759)
+ # TODO(https://github.com/apache/beam/issues/20759)
'pandas.core.frame.DataFrame.set_index': [
"df.set_index([s, s**2])",
],
-
'pandas.core.frame.DataFrame.set_axis': [
"df.set_axis(range(0,2), axis='index')",
],
- # TODO(https://github.com/apache/beam/issues/21014)
+ # TODO(https://github.com/apache/beam/issues/21014)
'pandas.core.frame.DataFrame.value_counts': [
- 'df.value_counts(dropna=False)'
+ 'df.value_counts(dropna=False)'
],
-
'pandas.core.frame.DataFrame.to_timestamp': ['*']
},
skip={
@@ -315,14 +309,12 @@ class DoctestTest(unittest.TestCase):
'*': [
# mul doesn't work in Beam with axis='index'.
"df.mul({'circle': 0, 'triangle': 2, 'rectangle': 3}, "
- "axis='index')",
- # eq doesn't work with axis='index'.
+ "axis='index')", # eq doesn't work with axis='index'.
"df.eq([250, 250, 100], axis='index')",
# New test in Pandas 2.1 that uses indexes.
'df != pd.Series([100, 250], index=["cost", "revenue"])',
# New test in Pandas 2.1 that uses indexes.
'df.le(df_multindex, level=1)'
-
],
# DeferredDataFrame doesn't implement the DF interchange protocol.
'pandas.core.frame.DataFrame.__dataframe__': ['*'],
@@ -335,20 +327,17 @@ class DoctestTest(unittest.TestCase):
'df',
'df2 = pd.DataFrame(data=df1, index=["a", "c"])',
'df2',
- ],
- # s2 created with reindex
+ ], # s2 created with reindex
'pandas.core.frame.DataFrame.dot': [
'df.dot(s2)',
],
-
'pandas.core.frame.DataFrame.resample': ['df'],
'pandas.core.frame.DataFrame.asfreq': ['*'],
# Throws NotImplementedError when modifying df
'pandas.core.frame.DataFrame.axes': [
# Returns deferred index.
'df.axes',
- ],
- # Skipped because the relies on loc to set cells in df2
+ ], # Skipped because the relies on loc to set cells in df2
'pandas.core.frame.DataFrame.compare': ['*'],
'pandas.core.frame.DataFrame.cov': [
# Relies on setting entries ahead of time.
@@ -371,8 +360,7 @@ class DoctestTest(unittest.TestCase):
# This should pass as set_axis(axis='columns')
# and fail with set_axis(axis='index')
"df.set_axis(['a', 'b', 'c'], axis='index')"
- ],
- # Beam's implementation takes a filepath as an argument.
+ ], # Beam's implementation takes a filepath as an argument.
'pandas.core.frame.DataFrame.to_html': ['*'],
'pandas.core.frame.DataFrame.to_markdown': ['*'],
'pandas.core.frame.DataFrame.to_parquet': ['*'],
@@ -384,11 +372,10 @@ class DoctestTest(unittest.TestCase):
'df.insert(1, "newcol", [99, 99])',
'df.insert(0, "col1", [100, 100], allow_duplicates=True)'
],
-
'pandas.core.frame.DataFrame.to_records': [
'df.index = df.index.rename("I")',
- 'index_dtypes = f"<S{df.index.str.len().max()}"', # 1.x
- 'index_dtypes = "<S{}".format(df.index.str.len().max())', #0.x
+ 'index_dtypes = f"<S{df.index.str.len().max()}"', # 1.x
+ 'index_dtypes = "<S{}".format(df.index.str.len().max())', #0.x
'df.to_records(index_dtypes=index_dtypes)',
],
# These tests use the static method pd.pivot_table, which doesn't
@@ -404,8 +391,7 @@ class DoctestTest(unittest.TestCase):
'df.pivot(index=["lev1", "lev2"], columns=["lev3"],values="values")',
'df.pivot(index="lev1", columns=["lev2", "lev3"], values="values")',
'df.pivot(index=["lev1", "lev2"], columns=["lev3"], values="values")',
- ],
- # pylint: enable=line-too-long
+ ], # pylint: enable=line-too-long
'pandas.core.frame.DataFrame.append': [
'df',
# pylint: disable=line-too-long
@@ -416,14 +402,12 @@ class DoctestTest(unittest.TestCase):
# Fails when result is a singleton:
# https://github.com/apache/beam/issues/28559
'pandas.core.frame.DataFrame.kurt': [
- 'df.kurt(axis=None).round(6)',
- 's.kurt()'
+ 'df.kurt(axis=None).round(6)', 's.kurt()'
],
# Fails when result is a singleton:
# https://github.com/apache/beam/issues/28559
'pandas.core.frame.DataFrame.sem': [
- 'df.sem().round(6)',
- 's.sem().round(6)'
+ 'df.sem().round(6)', 's.sem().round(6)'
],
'pandas.core.frame.DataFrame.melt': [
"df.columns = [list('ABC'), list('DEF')]", "df"
@@ -458,8 +442,7 @@ class DoctestTest(unittest.TestCase):
'pandas.core.frame.DataFrame.map': [
'df_copy.iloc[0, 0] = pd.NA',
"df_copy.map(lambda x: len(str(x)), na_action='ignore')",
- ],
- # Skipped so we don't need to install natsort
+ ], # Skipped so we don't need to install natsort
'pandas.core.frame.DataFrame.sort_values': [
'from natsort import index_natsorted',
'df.sort_values(\n'
@@ -531,8 +514,7 @@ class DoctestTest(unittest.TestCase):
's1.append(s2, ignore_index=True)',
],
'pandas.core.series.Series.replace': [
- "s.replace([1, 2], method='bfill')",
- # Relies on method='pad'
+ "s.replace([1, 2], method='bfill')", # Relies on method='pad'
"s.replace('a')",
# Relies on method='pad'
# value=None is not valid for pandas < 1.4
@@ -581,16 +563,14 @@ class DoctestTest(unittest.TestCase):
# indexes are identical and elide the alignment.
'ser.groupby(ser > 100).mean()',
],
- 'pandas.core.series.Series.asfreq': ['*'],
- # error formatting
+ 'pandas.core.series.Series.asfreq': ['*'], # error formatting
'pandas.core.series.Series.append': [
's1.append(s2, verify_integrity=True)',
],
'pandas.core.series.Series.cov': [
# Differs in LSB on jenkins.
"s1.cov(s2)",
- ],
- # Test framework doesn't materialze DeferredIndex.
+ ], # Test framework doesn't materialze DeferredIndex.
'pandas.core.series.Series.keys': ['s.keys()'],
# Skipped idxmax/idxmin due an issue with the test framework
'pandas.core.series.Series.idxmin': ['s.idxmin()'],
@@ -620,14 +600,12 @@ class DoctestTest(unittest.TestCase):
# Fails when result is a singleton:
# https://github.com/apache/beam/issues/28559
'pandas.core.series.Series.kurt': [
- 'df.kurt(axis=None).round(6)',
- 's.kurt()'
+ 'df.kurt(axis=None).round(6)', 's.kurt()'
],
# Fails when result is a singleton:
# https://github.com/apache/beam/issues/28559
'pandas.core.series.Series.sem': [
- 'df.sem().round(6)',
- 's.sem().round(6)'
+ 'df.sem().round(6)', 's.sem().round(6)'
],
})
self.assertEqual(result.failed, 0)
@@ -675,13 +653,13 @@ class DoctestTest(unittest.TestCase):
"pd.Series(['foo', 'fuz', np.nan]).str.replace('f', repr)"
],
- # output has incorrect formatting in 1.2.x
+ # output has incorrect formatting in 1.2.x
f'{module_name}.StringMethods.extractall': ['*'],
# For split and rsplit, if expand=True, then the series
# must be of CategoricalDtype, which pandas doesn't convert to
f'{module_name}.StringMethods.rsplit': [
- 's.str.split(r"\\+|=", expand=True)', # for pandas<1.4
+ 's.str.split(r"\\+|=", expand=True)', # for pandas<1.4
's.str.split(expand=True)',
's.str.rsplit("/", n=1, expand=True)',
's.str.split(r"and|plus", expand=True)',
@@ -692,7 +670,7 @@ class DoctestTest(unittest.TestCase):
's.str.split(r"\\.jpg", regex=False, expand=True)'
],
f'{module_name}.StringMethods.split': [
- 's.str.split(r"\\+|=", expand=True)', # for pandas<1.4
+ 's.str.split(r"\\+|=", expand=True)', # for pandas<1.4
's.str.split(expand=True)',
's.str.rsplit("/", n=1, expand=True)',
's.str.split(r"and|plus", expand=True)',
@@ -741,16 +719,16 @@ class DoctestTest(unittest.TestCase):
],
'pandas.core.indexes.accessors.TimedeltaProperties.to_pytimedelta': [
'*'
- ],
- # pylint: enable=line-too-long
- # Test uses to_datetime. Beam calls to_datetime element-wise, and
- # therefore the .tz attribute is not evaluated on entire Series.
- # Hence, .tz becomes None, unless explicitly set.
- # See: see test_tz_with_utc_zone_set_explicitly
+ ], # pylint: enable=line-too-long
+ # Test uses to_datetime. Beam calls to_datetime element-wise, and
+ # therefore the .tz attribute is not evaluated on entire Series.
+ # Hence, .tz becomes None, unless explicitly set.
+ # See: see test_tz_with_utc_zone_set_explicitly
'pandas.core.indexes.accessors.DatetimeProperties.tz': ['*'],
})
datetimelike_result = doctests.testmod(
- pd.core.arrays.datetimelike, use_beam=False,
+ pd.core.arrays.datetimelike,
+ use_beam=False,
not_implemented_ok={
# Beam Dataframes don't implement a deferred to_timedelta operation.
# Top-level issue: https://github.com/apache/beam/issues/20318
@@ -758,14 +736,12 @@ class DoctestTest(unittest.TestCase):
"ser = pd.Series(pd.to_timedelta([1, 2, 3], unit='d'))",
"tdelta_idx = pd.to_timedelta([1, 2, 3], unit='D')",
'tdelta_idx = pd.to_timedelta(["0 days", "10 days", "20 days"])', # pylint: disable=line-too-long
-
"tdelta_idx",
"tdelta_idx.inferred_freq",
"tdelta_idx.mean()",
],
})
-
datetime_result = doctests.testmod(
pd.core.arrays.datetimes,
use_beam=False,
@@ -782,7 +758,8 @@ class DoctestTest(unittest.TestCase):
'*': [
"ser = pd.Series(pd.to_timedelta([1, 2, 3], unit='d'))",
"tdelta_idx = pd.to_timedelta([1, 2, 3], unit='D')",
- 'tdelta_idx = pd.to_timedelta(["0 days", "10 days", "20 days"])'], # pylint: disable=line-too-long
+ 'tdelta_idx = pd.to_timedelta(["0 days", "10 days", "20 days"])'
+ ], # pylint: disable=line-too-long
# Verifies index version of this method
'pandas.core.arrays.datetimes.DatetimeArray.to_period': [
'df.index.to_period("M")'
@@ -872,8 +849,7 @@ class DoctestTest(unittest.TestCase):
'pandas.core.groupby.groupby.GroupBy.resample': [
'df.iloc[2, 0] = 5',
'df',
- ],
- # df is reassigned
+ ], # df is reassigned
'pandas.core.groupby.groupby.GroupBy.rank': ['df'],
# TODO: Raise wont implement for list passed as a grouping column
# Currently raises unhashable type: list
@@ -887,11 +863,10 @@ class DoctestTest(unittest.TestCase):
pd.core.groupby.generic,
use_beam=False,
wont_implement_ok={
- '*' : [
+ '*': [
# resample is WontImpl.
"ser.resample('MS').nunique()",
- ],
- # TODO: Is take actually deprecated?
+ ], # TODO: Is take actually deprecated?
'pandas.core.groupby.generic.DataFrameGroupBy.take': ['*'],
'pandas.core.groupby.generic.SeriesGroupBy.take': ['*'],
'pandas.core.groupby.generic.SeriesGroupBy.nsmallest': [
@@ -945,23 +920,24 @@ class DoctestTest(unittest.TestCase):
"df.loc[df.index[:5], 'a'] = np.nan",
"df.loc[df.index[5:10], 'b'] = np.nan",
"df.cov(min_periods=12)",
- ],
- # These examples rely on grouping by a list
+ ], # These examples rely on grouping by a list
'pandas.core.groupby.generic.SeriesGroupBy.aggregate': ['*'],
'pandas.core.groupby.generic.DataFrameGroupBy.aggregate': ['*'],
# Skipped idxmax/idxmin due an issue with the test framework
'pandas.core.groupby.generic.SeriesGroupBy.idxmin': ['s.idxmin()'],
'pandas.core.groupby.generic.SeriesGroupBy.idxmax': ['s.idxmax()'],
# Order-sensitive operations. TODO: Return a better error message.
- 'pandas.core.groupby.generic.SeriesGroupBy.is_monotonic_increasing': ['*'], # pylint: disable=line-too-long
- 'pandas.core.groupby.generic.SeriesGroupBy.is_monotonic_decreasing': ['*'], # pylint: disable=line-too-long
+ 'pandas.core.groupby.generic.SeriesGroupBy.is_monotonic_increasing': [
+ '*'
+ ], # pylint: disable=line-too-long
+ 'pandas.core.groupby.generic.SeriesGroupBy.is_monotonic_decreasing': [
+ '*'
+ ], # pylint: disable=line-too-long
# Uses as_index, which is currently not_implemented
'pandas.core.groupby.generic.DataFrameGroupBy.value_counts': [
- "df.groupby('gender', as_index=False).value_counts()",
- # pylint: disable=line-too-long
+ "df.groupby('gender', as_index=False).value_counts()", # pylint: disable=line-too-long
"df.groupby('gender', as_index=False).value_counts(normalize=True)",
- ],
- # These examples rely on grouping by a list
+ ], # These examples rely on grouping by a list
'pandas.core.groupby.generic.SeriesGroupBy.fillna': ['*'],
# These examples rely on grouping by a list
'pandas.core.groupby.generic.DataFrameGroupBy.fillna': ['*'],
@@ -972,8 +948,7 @@ class DoctestTest(unittest.TestCase):
# Named aggregation not supported yet.
'pandas.core.groupby.generic.NamedAgg': [
'df.groupby("key").agg(result_a=agg_a, result_1=agg_1)'
- ],
- # These examples rely on grouping by a list
+ ], # These examples rely on grouping by a list
'pandas.core.groupby.generic.DataFrameGroupBy.transform': ['*'],
# These examples rely on grouping by a list
'pandas.core.groupby.generic.SeriesGroupBy.transform': ['*'],
@@ -1024,7 +999,9 @@ class DoctestTest(unittest.TestCase):
'pivot': ['*'],
'to_datetime': ['s.head()'],
'to_pickle': ['*'],
- 'unique': ['pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values)'], # pylint: disable=line-too-long
+ 'unique': [
+ 'pd.unique(pd.Series([("a", "b"), ("b", "a"), ("a", "c"), ("b", "a")]).values)'
+ ], # pylint: disable=line-too-long
'melt': [
"pd.melt(df, id_vars=['A'], value_vars=['B'])",
"pd.melt(df, id_vars=['A'], value_vars=['B', 'C'])",
@@ -1039,34 +1016,27 @@ class DoctestTest(unittest.TestCase):
'concat': [
'pd.concat([df5, df6], verify_integrity=True)',
'pd.concat([df7, new_row.to_frame().T], ignore_index=True)'
- ],
- # doctest DeprecationWarning
- 'melt': ['df'],
- # Order-sensitive re-indexing.
+ ], # doctest DeprecationWarning
+ 'melt': ['df'], # Order-sensitive re-indexing.
'merge': [
"df1.merge(df2, left_on='lkey', right_on='rkey')",
"df1.merge(df2, left_on='lkey', right_on='rkey',\n"
" suffixes=('_left', '_right'))",
"df1.merge(df2, how='left', on='a')",
- ],
- # Not an actual test.
+ ], # Not an actual test.
'option_context': ['*'],
'factorize': ['codes', 'uniques'],
# Bad top-level use of un-imported function.
'merge_ordered': [
'merge_ordered(df1, df2, fill_method="ffill", left_by="group")'
- ],
- # Expected error.
+ ], # Expected error.
'pivot': [
"df.pivot(index='foo', columns='bar', values='baz')",
"df.pivot(index='foo', columns='bar')['baz']",
- "df.pivot(index='foo', columns='bar', values=['baz', 'zoo'])",
- # pylint: disable=line-too-long
- 'df.pivot(index="lev1", columns=["lev2", "lev3"],values="values")',
- # pylint: disable=line-too-long
+ "df.pivot(index='foo', columns='bar', values=['baz', 'zoo'])", # pylint: disable=line-too-long
+ 'df.pivot(index="lev1", columns=["lev2", "lev3"],values="values")', # pylint: disable=line-too-long
'df.pivot(index=["lev1", "lev2"], columns=["lev3"],values="values")'
- ],
- # Never written.
+ ], # Never written.
'to_pickle': ['os.remove("./dummy.pkl")'],
**skip_reads
})
diff --git a/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py b/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py
index a8139675ad3..061529f5aea 100644
--- a/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py
+++ b/sdks/python/apache_beam/dataframe/pandas_top_level_functions.py
@@ -28,6 +28,7 @@ from apache_beam.dataframe import partitionings
def _call_on_first_arg(name):
+
def wrapper(target, *args, **kwargs):
if isinstance(target, frame_base.DeferredBase):
return getattr(target, name)(*args, **kwargs)
@@ -131,14 +132,9 @@ class DeferredPandasModule(object):
expressions.ComputedExpression(
'concat',
lambda *objs: pd.concat(
- objs,
- axis=axis,
- join=join,
- ignore_index=ignore_index,
- keys=keys,
- levels=levels,
- names=names,
- verify_integrity=verify_integrity), # yapf break
+ objs, axis=axis, join=join, ignore_index=ignore_index, keys=
+ keys, levels=levels, names=names, verify_integrity=
+ verify_integrity), # yapf break
exprs,
requires_partition_by=required_partitioning,
preserves_partition_by=preserves_partitioning))
diff --git a/sdks/python/apache_beam/dataframe/partitionings.py b/sdks/python/apache_beam/dataframe/partitionings.py
index 1fe760fe858..09409908928 100644
--- a/sdks/python/apache_beam/dataframe/partitionings.py
+++ b/sdks/python/apache_beam/dataframe/partitionings.py
@@ -28,6 +28,7 @@ Frame = TypeVar('Frame', bound=pd.core.generic.NDFrame)
class Partitioning(object):
"""A class representing a (consistent) partitioning of dataframe objects.
"""
+
def __repr__(self):
return self.__class__.__name__
@@ -71,6 +72,7 @@ class Index(Partitioning):
The ordering is implemented via the is_subpartitioning_of method, where the
examples on the right are subpartitionings of the examples on the left above.
"""
+
def __init__(self, levels=None):
self._levels = levels
@@ -148,6 +150,7 @@ class Index(Partitioning):
class Singleton(Partitioning):
"""A partitioning of all the data into a single partition.
"""
+
def __init__(self, reason=None):
self._reason = reason
@@ -189,6 +192,7 @@ class JoinIndex(Partitioning):
Expressions desiring to make use of this index should simply declare a
requirement of JoinIndex().
"""
+
def __init__(self, ancestor=None):
self._ancestor = ancestor
@@ -229,6 +233,7 @@ class JoinIndex(Partitioning):
class Arbitrary(Partitioning):
"""A partitioning imposing no constraints on the actual partitioning.
"""
+
def __eq__(self, other):
return type(self) == type(other)
diff --git a/sdks/python/apache_beam/dataframe/schemas.py b/sdks/python/apache_beam/dataframe/schemas.py
index f849ab11e77..c45cee5e04b 100644
--- a/sdks/python/apache_beam/dataframe/schemas.py
+++ b/sdks/python/apache_beam/dataframe/schemas.py
@@ -61,6 +61,7 @@ class BatchRowsAsDataFrame(beam.PTransform):
Batching parameters are inherited from
:class:`~apache_beam.transforms.util.BatchElements`.
"""
+
def __init__(self, *args, proxy=None, **kwargs):
self._batch_elements_transform = BatchElements(*args, **kwargs)
self._proxy = proxy
@@ -204,6 +205,7 @@ class UnbatchPandas(beam.PTransform):
levels are unnamed (name=None), or if any of the names are not unique
among all column and index names.
"""
+
def __init__(self, proxy, include_indexes=False):
self._proxy = proxy
self._include_indexes = include_indexes
diff --git a/sdks/python/apache_beam/dataframe/schemas_test.py b/sdks/python/apache_beam/dataframe/schemas_test.py
index 4c196e29e71..634443e4bf8 100644
--- a/sdks/python/apache_beam/dataframe/schemas_test.py
+++ b/sdks/python/apache_beam/dataframe/schemas_test.py
@@ -47,6 +47,7 @@ coders_registry.register_coder(Animal, RowCoder)
def matches_df(expected):
+
def check_df_pcoll_equal(actual):
actual = pd.concat(actual)
sorted_actual = actual.sort_values(by=list(actual.columns)).reset_index(
@@ -157,6 +158,7 @@ def test_name_func(testcase_func, param_num, params):
class SchemasTest(unittest.TestCase):
+
def test_simple_df(self):
expected = pd.DataFrame({
'name': list(str(i) for i in range(5)),
@@ -265,6 +267,7 @@ class SchemasTest(unittest.TestCase):
assert_that(res, equal_to([('Falcon', 375.), ('Parrot', 25.)]))
def assert_typehints_equal(self, left, right):
+
def maybe_drop_rowtypeconstraint(typehint):
if isinstance(typehint, row_type.RowTypeConstraint):
return typehint.user_type
diff --git a/sdks/python/apache_beam/dataframe/transforms.py b/sdks/python/apache_beam/dataframe/transforms.py
index 59e5eec05d2..354f18618f5 100644
--- a/sdks/python/apache_beam/dataframe/transforms.py
+++ b/sdks/python/apache_beam/dataframe/transforms.py
@@ -90,6 +90,7 @@ class DataframeTransform(transforms.PTransform):
.. _schema-aware:
https://beam.apache.org/documentation/programming-guide/#what-is-a-schema
"""
+
def __init__(
self, func, proxy=None, yield_elements="schemas", include_indexes=False):
self._func = func
@@ -138,6 +139,7 @@ class DataframeTransform(transforms.PTransform):
class _DataframeExpressionsTransform(transforms.PTransform):
+
def __init__(self, outputs):
self._outputs = outputs
@@ -163,9 +165,11 @@ class _DataframeExpressionsTransform(transforms.PTransform):
Logically, `_apply_deferred_ops({x: a, y: b}, {f: F(x, y), g: G(x, y)})`
returns `{f: F(a, b), g: G(a, b)}`.
"""
+
class ComputeStage(beam.PTransform):
"""A helper transform that computes a single stage of operations.
"""
+
def __init__(self, stage):
self.stage = stage
@@ -208,12 +212,13 @@ class _DataframeExpressionsTransform(transforms.PTransform):
| 'SumSizes' >> beam.CombineGlobally(sum)
| 'NumPartitions' >> beam.Map(
lambda size: max(
- MIN_PARTITIONS,
- min(MAX_PARTITIONS, size // TARGET_PARTITION_SIZE))))
+ MIN_PARTITIONS, min(
+ MAX_PARTITIONS, size // TARGET_PARTITION_SIZE))))
partition_fn = self.stage.partitioning.partition_fn
class Partition(beam.PTransform):
+
def expand(self, pcoll):
return (
pcoll
@@ -245,10 +250,11 @@ class _DataframeExpressionsTransform(transforms.PTransform):
# Actually evaluate the expressions.
def evaluate(partition, stage=self.stage, **side_inputs):
+
def lookup(expr):
# Use proxy if there's no data in this partition
- return expr.proxy(
- ).iloc[:0] if partition[expr._id] is None else partition[expr._id]
+ return expr.proxy().iloc[:0] if partition[
+ expr._id] is None else partition[expr._id]
session = expressions.Session(
dict([(expr, lookup(expr)) for expr in tabular_inputs] +
@@ -265,6 +271,7 @@ class _DataframeExpressionsTransform(transforms.PTransform):
Note that these Dataframe "stages" contain a CoGBK and hence are often
split across multiple "executable" stages.
"""
+
def __init__(self, inputs, partitioning):
self.inputs = set(inputs)
if (len(self.inputs) > 1 and
@@ -301,6 +308,7 @@ class _DataframeExpressionsTransform(transforms.PTransform):
"""Return the output partitioning of expr when computed in stage,
or returns None if the expression cannot be computed in this stage.
"""
+
def maybe_upgrade_to_join_index(partitioning):
if partitioning.is_subpartitioning_of(partitionings.JoinIndex()):
return partitionings.JoinIndex(expr)
@@ -420,8 +428,10 @@ class _DataframeExpressionsTransform(transforms.PTransform):
@_memoize
def stage_to_result(stage):
- return {expr._id: expr_to_pcoll(expr)
- for expr in stage.inputs} | ComputeStage(stage)
+ return {
+ expr._id: expr_to_pcoll(expr)
+ for expr in stage.inputs
+ } | ComputeStage(stage)
@_memoize
def expr_to_pcoll(expr):
@@ -484,6 +494,7 @@ def _total_memory_usage(frame):
class _PreBatch(beam.DoFn):
+
def __init__(
self, target_size=TARGET_PARTITION_SIZE, min_size=MIN_PARTITION_SIZE):
self._target_size = target_size
@@ -519,6 +530,7 @@ class _ReBatch(beam.DoFn):
Also groups across partitions, up to a given data size, to recover some
efficiency in the face of over-partitioning.
"""
+
def __init__(
self, target_size=TARGET_PARTITION_SIZE, min_size=MIN_PARTITION_SIZE):
self._target_size = target_size
diff --git a/sdks/python/apache_beam/dataframe/transforms_test.py b/sdks/python/apache_beam/dataframe/transforms_test.py
index a143606cc91..25829a256fa 100644
--- a/sdks/python/apache_beam/dataframe/transforms_test.py
+++ b/sdks/python/apache_beam/dataframe/transforms_test.py
@@ -74,6 +74,7 @@ coders.registry.register_coder(Nested, coders.RowCoder)
class TransformTest(unittest.TestCase):
+
def run_scenario(self, input, func):
expected = func(input)
@@ -274,6 +275,7 @@ class TransformTest(unittest.TestCase):
proxy = one_series[:0]
def equal_to_series(expected):
+
def check(actual):
actual = pd.concat(actual)
if not expected.equals(actual):
@@ -302,8 +304,7 @@ class TransformTest(unittest.TestCase):
assert_that(
dict(x=one, y=two)
| 'DictIn' >> transforms.DataframeTransform(
- lambda x,
- y: (x + y),
+ lambda x, y: (x + y),
proxy=dict(x=proxy, y=proxy),
yield_elements='pandas'),
equal_to_series(three_series),
@@ -348,14 +349,14 @@ class TransformTest(unittest.TestCase):
with expressions.allow_non_parallel_operations():
self.run_scenario(
- df,
- lambda df: df.rename(
+ df, lambda df: df.rename(
columns={'B': 'C'}, index={
0: 2, 2: 0
}, errors='raise'))
class FusionTest(unittest.TestCase):
+
@staticmethod
def fused_stages(p):
return p.result.monitoring_metrics().query(
@@ -381,6 +382,7 @@ class FusionTest(unittest.TestCase):
self.assertEqual(self.fused_stages(p), 1)
def test_column_manipulation(self):
+
def set_column(df, name, s):
df[name] = s
return df
@@ -394,6 +396,7 @@ class FusionTest(unittest.TestCase):
class TransformPartsTest(unittest.TestCase):
+
def test_rebatch(self):
with beam.Pipeline() as p:
sA = pd.Series(range(1000))
diff --git a/sdks/python/apache_beam/examples/avro_nyc_trips.py b/sdks/python/apache_beam/examples/avro_nyc_trips.py
index 23d25649dad..0f7cfb301c8 100644
--- a/sdks/python/apache_beam/examples/avro_nyc_trips.py
+++ b/sdks/python/apache_beam/examples/avro_nyc_trips.py
@@ -125,6 +125,7 @@ class CreateKeyWithServiceAndDay(beam.DoFn):
company name and the day of the week of the ride. The value is the original
record dictionary object.
"""
+
def process(self, record: dict):
options = {
'HV0002': 'Juno', 'HV0003': 'Uber', 'HV0004': 'Via', 'HV0005': 'Lyft'
@@ -154,6 +155,7 @@ class CalculatePricePerAttribute(beam.CombineFn):
hire vehicle service. And calculates the price per mile, minute, and trip
for both the driver and passenger.
"""
+
def create_accumulator(self):
total_price = 0.0
total_driver_pay = 0.0
@@ -180,13 +182,9 @@ class CalculatePricePerAttribute(beam.CombineFn):
return (
total_price + sum(
record[name] for name in (
- 'base_passenger_fare',
- 'tolls',
- 'bcf',
- 'sales_tax',
- 'congestion_surcharge',
- 'airport_fee',
- 'tips') if record[name] is not None),
+ 'base_passenger_fare', 'tolls', 'bcf', 'sales_tax',
+ 'congestion_surcharge', 'airport_fee', 'tips')
+ if record[name] is not None),
total_driver_pay + record['driver_pay'] + record['tips'],
total_trip_miles + record['trip_miles'],
total_trip_time + record['trip_time'],
diff --git a/sdks/python/apache_beam/examples/avro_nyc_trips_test.py b/sdks/python/apache_beam/examples/avro_nyc_trips_test.py
index 0dd31962a2f..67141f23c4c 100644
--- a/sdks/python/apache_beam/examples/avro_nyc_trips_test.py
+++ b/sdks/python/apache_beam/examples/avro_nyc_trips_test.py
@@ -29,6 +29,7 @@ from apache_beam.testing.util import equal_to
class AvroNycTripsTest(unittest.TestCase):
+
def test_create_key_with_service_and_day(self):
RECORDS = [
{
diff --git a/sdks/python/apache_beam/examples/complete/autocomplete.py b/sdks/python/apache_beam/examples/complete/autocomplete.py
index 4e4c5143b96..e4e45dd1a7c 100644
--- a/sdks/python/apache_beam/examples/complete/autocomplete.py
+++ b/sdks/python/apache_beam/examples/complete/autocomplete.py
@@ -57,6 +57,7 @@ def run(argv=None):
class TopPerPrefix(beam.PTransform):
+
def __init__(self, count):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
diff --git a/sdks/python/apache_beam/examples/complete/autocomplete_it_test.py b/sdks/python/apache_beam/examples/complete/autocomplete_it_test.py
index a19af587318..101eccc84e3 100644
--- a/sdks/python/apache_beam/examples/complete/autocomplete_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/autocomplete_it_test.py
@@ -32,6 +32,7 @@ from apache_beam.testing.test_utils import read_files_from_pattern
def format_output_file(output_string):
+
def extract_prefix_topk_words_tuples(line):
match = re.match(r'(.*): \[(.*)\]', line)
prefix = match.group(1)
diff --git a/sdks/python/apache_beam/examples/complete/distribopt.py b/sdks/python/apache_beam/examples/complete/distribopt.py
index 89c312fcbf5..3fd9a70691d 100644
--- a/sdks/python/apache_beam/examples/complete/distribopt.py
+++ b/sdks/python/apache_beam/examples/complete/distribopt.py
@@ -68,6 +68,7 @@ from apache_beam.options.pipeline_options import SetupOptions
class Simulator(object):
"""Greenhouse simulation for the optimization of greenhouse parameters."""
+
def __init__(self, quantities):
self.quantities = np.atleast_1d(quantities)
@@ -99,6 +100,7 @@ class CreateGrid(beam.PTransform):
}
Output: tuple (mapping_identifier, {crop -> greenhouse})
"""
+
class PreGenerateMappings(beam.DoFn):
"""ParDo implementation forming based on two elements a small sub grid.
@@ -107,6 +109,7 @@ class CreateGrid(beam.PTransform):
two tuples, and a list of remaining records. Both serve as an input to
GenerateMappings.
"""
+
def process(self, element):
records = list(element[1])
# Split of 2 crops and pre-generate the subgrid.
@@ -135,6 +138,7 @@ class CreateGrid(beam.PTransform):
Input: output of PreGenerateMappings
Output: tuples of the form (mapping_identifier, {crop -> greenhouse})
"""
+
@staticmethod
def _coordinates_to_greenhouse(coordinates, greenhouses, crops):
# Map the grid coordinates back to greenhouse labels
@@ -185,6 +189,7 @@ class CreateGrid(beam.PTransform):
class OptimizeGrid(beam.PTransform):
"""A transform for optimizing all greenhouses of the mapping grid."""
+
class CreateOptimizationTasks(beam.DoFn):
"""
Create tasks for optimization.
@@ -192,6 +197,7 @@ class OptimizeGrid(beam.PTransform):
Input: (mapping_identifier, {crop -> greenhouse})
Output: ((mapping_identifier, greenhouse), [(crop, quantity),...])
"""
+
def process(self, element, quantities):
mapping_identifier, mapping = element
@@ -213,6 +219,7 @@ class OptimizeGrid(beam.PTransform):
- solution: (mapping_identifier, (greenhouse, [production parameters]))
- costs: (crop, greenhouse, mapping_identifier, cost)
"""
+
@staticmethod
def _optimize_production_parameters(sim):
# setup initial starting point & bounds
@@ -247,6 +254,7 @@ class OptimizeGrid(beam.PTransform):
class CreateTransportData(beam.DoFn):
"""Transform records to pvalues ((crop, greenhouse), transport_cost)"""
+
def process(self, record):
crop = record['crop']
for greenhouse, transport_cost in record['transport_costs']:
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi.py b/sdks/python/apache_beam/examples/complete/estimate_pi.py
index 530a270308d..350f078b138 100644
--- a/sdks/python/apache_beam/examples/complete/estimate_pi.py
+++ b/sdks/python/apache_beam/examples/complete/estimate_pi.py
@@ -81,12 +81,14 @@ def combine_results(results):
class JsonCoder(object):
"""A JSON coder used to format the final result."""
+
def encode(self, x):
return json.dumps(x).encode('utf-8')
class EstimatePiTransform(beam.PTransform):
"""Runs 10M trials, and combine the results to estimate pi."""
+
def __init__(self, tries_per_work_item=100000):
self.tries_per_work_item = tries_per_work_item
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi_it_test.py b/sdks/python/apache_beam/examples/complete/estimate_pi_it_test.py
index bf6f8fc76c1..6365e1491db 100644
--- a/sdks/python/apache_beam/examples/complete/estimate_pi_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/estimate_pi_it_test.py
@@ -31,6 +31,7 @@ from apache_beam.testing.test_utils import read_files_from_pattern
class EstimatePiIT(unittest.TestCase):
+
@pytest.mark.no_xdist
@pytest.mark.examples_postcommit
def test_estimate_pi_output_file(self):
diff --git a/sdks/python/apache_beam/examples/complete/estimate_pi_test.py b/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
index ff224f4a9ca..909c4ab66b2 100644
--- a/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
+++ b/sdks/python/apache_beam/examples/complete/estimate_pi_test.py
@@ -29,6 +29,7 @@ from apache_beam.testing.util import assert_that
def in_between(lower, upper):
+
def _in_between(actual):
_, _, estimate = actual[0]
if estimate < lower or estimate > upper:
@@ -39,6 +40,7 @@ def in_between(lower, upper):
class EstimatePiTest(unittest.TestCase):
+
def test_basics(self):
with TestPipeline() as p:
result = p | 'Estimate' >> estimate_pi.EstimatePiTransform(5000)
diff --git a/sdks/python/apache_beam/examples/complete/game/game_stats.py b/sdks/python/apache_beam/examples/complete/game/game_stats.py
index 233d22b7542..95000677696 100644
--- a/sdks/python/apache_beam/examples/complete/game/game_stats.py
+++ b/sdks/python/apache_beam/examples/complete/game/game_stats.py
@@ -104,6 +104,7 @@ class ParseGameEventFn(beam.DoFn):
The human-readable time string is not used here.
"""
+
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -130,6 +131,7 @@ class ExtractAndSumScore(beam.PTransform):
The constructor argument `field` determines whether 'team' or 'user' info is
extracted.
"""
+
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -150,6 +152,7 @@ class TeamScoresDict(beam.DoFn):
formats everything together into a dictionary. The dictionary is in the format
{'bigquery_column': value}
"""
+
def process(self, team_score, window=beam.DoFn.WindowParam):
team, score = team_score
start = timestamp2str(int(window.start))
@@ -163,6 +166,7 @@ class TeamScoresDict(beam.DoFn):
class WriteToBigQuery(beam.PTransform):
"""Generate, format, and write BigQuery table row information."""
+
def __init__(self, table_name, dataset, schema, project):
"""Initializes the transform.
Args:
@@ -231,6 +235,7 @@ class CalculateSpammyUsers(beam.PTransform):
class UserSessionActivity(beam.DoFn):
"""Calculate and output an element's session duration, in seconds."""
+
def process(self, elem, window=beam.DoFn.WindowParam):
yield (window.end.micros - window.start.micros) // 1000000
diff --git a/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py b/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
index 48a105af527..21d00fad850 100644
--- a/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
+++ b/sdks/python/apache_beam/examples/complete/game/hourly_team_score.py
@@ -104,6 +104,7 @@ class ParseGameEventFn(beam.DoFn):
The human-readable time string is not used here.
"""
+
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -130,6 +131,7 @@ class ExtractAndSumScore(beam.PTransform):
The constructor argument `field` determines whether 'team' or 'user' info is
extracted.
"""
+
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -150,6 +152,7 @@ class TeamScoresDict(beam.DoFn):
formats everything together into a dictionary. The dictionary is in the format
{'bigquery_column': value}
"""
+
def process(self, team_score, window=beam.DoFn.WindowParam):
team, score = team_score
start = timestamp2str(int(window.start))
@@ -163,6 +166,7 @@ class TeamScoresDict(beam.DoFn):
class WriteToBigQuery(beam.PTransform):
"""Generate, format, and write BigQuery table row information."""
+
def __init__(self, table_name, dataset, schema, project):
"""Initializes the transform.
Args:
@@ -195,6 +199,7 @@ class WriteToBigQuery(beam.PTransform):
# [START main]
class HourlyTeamScore(beam.PTransform):
+
def __init__(self, start_min, stop_min, window_duration):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
diff --git a/sdks/python/apache_beam/examples/complete/game/leader_board.py b/sdks/python/apache_beam/examples/complete/game/leader_board.py
index 308e1e1cf5c..71a28406bf7 100644
--- a/sdks/python/apache_beam/examples/complete/game/leader_board.py
+++ b/sdks/python/apache_beam/examples/complete/game/leader_board.py
@@ -113,6 +113,7 @@ class ParseGameEventFn(beam.DoFn):
The human-readable time string is not used here.
"""
+
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -139,6 +140,7 @@ class ExtractAndSumScore(beam.PTransform):
The constructor argument `field` determines whether 'team' or 'user' info is
extracted.
"""
+
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -159,6 +161,7 @@ class TeamScoresDict(beam.DoFn):
formats everything together into a dictionary. The dictionary is in the format
{'bigquery_column': value}
"""
+
def process(self, team_score, window=beam.DoFn.WindowParam):
team, score = team_score
start = timestamp2str(int(window.start))
@@ -172,6 +175,7 @@ class TeamScoresDict(beam.DoFn):
class WriteToBigQuery(beam.PTransform):
"""Generate, format, and write BigQuery table row information."""
+
def __init__(self, table_name, dataset, schema, project):
"""Initializes the transform.
Args:
@@ -209,6 +213,7 @@ class CalculateTeamScores(beam.PTransform):
Extract team/score pairs from the event stream, using hour-long windows by
default.
"""
+
def __init__(self, team_window_duration, allowed_lateness):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -241,6 +246,7 @@ class CalculateUserScores(beam.PTransform):
"""Extract user/score pairs from the event stream using processing time, via
global windowing. Get periodic updates on all users' running scores.
"""
+
def __init__(self, allowed_lateness):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
diff --git a/sdks/python/apache_beam/examples/complete/game/user_score.py b/sdks/python/apache_beam/examples/complete/game/user_score.py
index 03f0d00fc30..00056cff586 100644
--- a/sdks/python/apache_beam/examples/complete/game/user_score.py
+++ b/sdks/python/apache_beam/examples/complete/game/user_score.py
@@ -96,6 +96,7 @@ class ParseGameEventFn(beam.DoFn):
The human-readable time string is not used here.
"""
+
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -123,6 +124,7 @@ class ExtractAndSumScore(beam.PTransform):
The constructor argument `field` determines whether 'team' or 'user' info is
extracted.
"""
+
def __init__(self, field):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -140,6 +142,7 @@ class ExtractAndSumScore(beam.PTransform):
class UserScore(beam.PTransform):
+
def expand(self, pcoll):
return (
pcoll
diff --git a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py
index 4f98105c66f..c08b968e0ab 100644
--- a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py
+++ b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset.py
@@ -46,6 +46,7 @@ def get_julia_set_point_color(element, c, n, max_iterations):
def generate_julia_set_colors(pipeline, c, n, max_iterations):
"""Compute julia set coordinates for each point in our set."""
+
def point_set(n):
for x in range(n):
for y in range(n):
diff --git a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py
index 6416831f426..0a69d89edb1 100644
--- a/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py
+++ b/sdks/python/apache_beam/examples/complete/juliaset/juliaset/juliaset_test.py
@@ -33,6 +33,7 @@ from apache_beam.testing.util import open_shards
@pytest.mark.examples_postcommit
class JuliaSetTest(unittest.TestCase):
+
def setUp(self):
self.test_files = {}
self.test_files['output_coord_file_name'] = self.generate_temp_file()
diff --git a/sdks/python/apache_beam/examples/complete/juliaset/setup.py b/sdks/python/apache_beam/examples/complete/juliaset/setup.py
index c3a9fe04376..05ecbd8439e 100644
--- a/sdks/python/apache_beam/examples/complete/juliaset/setup.py
+++ b/sdks/python/apache_beam/examples/complete/juliaset/setup.py
@@ -81,6 +81,7 @@ CUSTOM_COMMANDS = [['echo', 'Custom command worked!']]
class CustomCommands(setuptools.Command):
"""A setuptools Command class able to run arbitrary commands."""
+
def initialize_options(self):
pass
diff --git a/sdks/python/apache_beam/examples/complete/tfidf.py b/sdks/python/apache_beam/examples/complete/tfidf.py
index d7829f9d1c7..651e9ec3298 100644
--- a/sdks/python/apache_beam/examples/complete/tfidf.py
+++ b/sdks/python/apache_beam/examples/complete/tfidf.py
@@ -54,6 +54,7 @@ class TfIdf(beam.PTransform):
the value is a piece of the document's content.
The output is mapping from terms to scores for each document URI.
"""
+
def expand(self, uri_to_content):
# Compute the total number of documents, and prepare a singleton
diff --git a/sdks/python/apache_beam/examples/complete/tfidf_it_test.py b/sdks/python/apache_beam/examples/complete/tfidf_it_test.py
index 3ecbd0c1eca..6985f288b2a 100644
--- a/sdks/python/apache_beam/examples/complete/tfidf_it_test.py
+++ b/sdks/python/apache_beam/examples/complete/tfidf_it_test.py
@@ -41,6 +41,7 @@ EXPECTED_LINE_RE = r'\(u?\'([a-z]*)\', \(\'.*([0-9]\.txt)\', (.*)\)\)'
class TfIdfIT(unittest.TestCase):
+
@pytest.mark.examples_postcommit
@pytest.mark.sickbay_flink
def test_basics(self):
diff --git a/sdks/python/apache_beam/examples/complete/tfidf_test.py b/sdks/python/apache_beam/examples/complete/tfidf_test.py
index 085b9e2dd18..6b31dd05f35 100644
--- a/sdks/python/apache_beam/examples/complete/tfidf_test.py
+++ b/sdks/python/apache_beam/examples/complete/tfidf_test.py
@@ -36,6 +36,7 @@ EXPECTED_RESULTS = set([
class TfIdfTest(unittest.TestCase):
+
def test_tfidf_transform(self):
with TestPipeline() as p:
diff --git a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
index 50b026edf24..e6c0bc7e01f 100644
--- a/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
+++ b/sdks/python/apache_beam/examples/complete/top_wikipedia_sessions.py
@@ -76,6 +76,7 @@ class ComputeSessions(beam.PTransform):
A session is defined as a string of edits where each is separated from the
next by less than an hour.
"""
+
def expand(self, pcoll):
return (
pcoll
@@ -86,6 +87,7 @@ class ComputeSessions(beam.PTransform):
class TopPerMonth(beam.PTransform):
"""Computes the longest session ending in each month."""
+
def expand(self, pcoll):
return (
pcoll
@@ -110,6 +112,7 @@ def format_output(element, window=beam.DoFn.WindowParam):
class ComputeTopSessions(beam.PTransform):
"""Computes the top user sessions for each month."""
+
def __init__(self, sampling_threshold):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py
index 6f40b14226b..f858ecc2965 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input.py
@@ -43,6 +43,7 @@ from apache_beam.pvalue import AsSingleton
def create_groups(group_ids, corpus, word, ignore_corpus, ignore_word):
"""Generate groups given the input PCollections."""
+
def attach_corpus_fn(group, corpus, ignore):
selected = None
len_corpus = len(corpus)
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
index ba9b61e342b..3274c4a0e9b 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_side_input_test.py
@@ -30,6 +30,7 @@ from apache_beam.testing.util import equal_to
class BigQuerySideInputTest(unittest.TestCase):
+
def test_create_groups(self):
with TestPipeline() as p:
diff --git a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
index e9f3f679ad8..a8763bf79cc 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigquery_tornadoes_test.py
@@ -30,6 +30,7 @@ from apache_beam.testing.util import equal_to
class BigQueryTornadoesTest(unittest.TestCase):
+
def test_basics(self):
with TestPipeline() as p:
rows = (
diff --git a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
index 0a8c55d17d3..a6a28b56ba0 100644
--- a/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/bigtableio_it_test.py
@@ -66,6 +66,7 @@ class GenerateTestRows(beam.PTransform):
Bigtable Table.
"""
+
def __init__(self, number, project_id=None, instance_id=None, table_id=None):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
diff --git a/sdks/python/apache_beam/examples/cookbook/coders.py b/sdks/python/apache_beam/examples/cookbook/coders.py
index 33b63fd2954..655f7a88d7c 100644
--- a/sdks/python/apache_beam/examples/cookbook/coders.py
+++ b/sdks/python/apache_beam/examples/cookbook/coders.py
@@ -44,6 +44,7 @@ from apache_beam.options.pipeline_options import SetupOptions
class JsonCoder(Coder):
"""A JSON coder interpreting each line as a JSON string."""
+
def encode(self, x):
return json.dumps(x).encode('utf-8')
diff --git a/sdks/python/apache_beam/examples/cookbook/coders_it_test.py b/sdks/python/apache_beam/examples/cookbook/coders_it_test.py
index 941311ce5dc..837c0c68423 100644
--- a/sdks/python/apache_beam/examples/cookbook/coders_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/coders_it_test.py
@@ -32,6 +32,7 @@ from apache_beam.testing.test_utils import read_files_from_pattern
def format_result(result_string):
+
def format_tuple(result_elem_list):
[country, counter] = result_elem_list
return country, int(counter.strip())
diff --git a/sdks/python/apache_beam/examples/cookbook/combiners_test.py b/sdks/python/apache_beam/examples/cookbook/combiners_test.py
index c24f60b5dcf..a20e7a02c73 100644
--- a/sdks/python/apache_beam/examples/cookbook/combiners_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/combiners_test.py
@@ -57,6 +57,7 @@ class CombinersTest(unittest.TestCase):
def test_combine_per_key_with_custom_callable(self):
"""CombinePerKey using a custom function reducing iterables."""
+
def multiply(values):
result = 1
for v in values:
diff --git a/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py b/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py
index a922216a522..51325ccab33 100644
--- a/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py
+++ b/sdks/python/apache_beam/examples/cookbook/custom_ptransform.py
@@ -36,6 +36,7 @@ from apache_beam.options.pipeline_options import PipelineOptions
class Count1(beam.PTransform):
"""Count as a subclass of PTransform, with an apply method."""
+
def expand(self, pcoll):
return (
pcoll
diff --git a/sdks/python/apache_beam/examples/cookbook/custom_ptransform_it_test.py b/sdks/python/apache_beam/examples/cookbook/custom_ptransform_it_test.py
index 9ad0c52bf23..52e4fdc77df 100644
--- a/sdks/python/apache_beam/examples/cookbook/custom_ptransform_it_test.py
+++ b/sdks/python/apache_beam/examples/cookbook/custom_ptransform_it_test.py
@@ -31,6 +31,7 @@ from apache_beam.testing.test_utils import read_files_from_pattern
def format_result(result_string):
+
def format_tuple(result_elem_list):
[country, counter] = result_elem_list
return country, int(counter.strip())
diff --git a/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py b/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py
index 9d71ac32aff..cc2db388587 100644
--- a/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py
+++ b/sdks/python/apache_beam/examples/cookbook/datastore_wordcount.py
@@ -81,6 +81,7 @@ from apache_beam.options.pipeline_options import PipelineOptions
@beam.typehints.with_output_types(Text)
class WordExtractingDoFn(beam.DoFn):
"""Parse each line of input text into words."""
+
def __init__(self):
self.empty_line_counter = Metrics.counter('main', 'empty_lines')
self.word_length_counter = Metrics.counter('main', 'word_lengths')
@@ -111,6 +112,7 @@ class WordExtractingDoFn(beam.DoFn):
class EntityWrapper(object):
"""Create a Cloud Datastore entity from the given string."""
+
def __init__(self, project, namespace, kind, ancestor):
self._project = project
self._namespace = namespace
diff --git a/sdks/python/apache_beam/examples/cookbook/group_with_coder.py b/sdks/python/apache_beam/examples/cookbook/group_with_coder.py
index 8a959138d3d..2ce443b85af 100644
--- a/sdks/python/apache_beam/examples/cookbook/group_with_coder.py
+++ b/sdks/python/apache_beam/examples/cookbook/group_with_coder.py
@@ -42,12 +42,14 @@ from apache_beam.typehints.decorators import with_output_types
class Player(object):
"""A custom class used as a key in combine/group transforms."""
+
def __init__(self, name):
self.name = name
class PlayerCoder(coders.Coder):
"""A custom coder for the Player class."""
+
def encode(self, o):
"""Encode to bytes with a trace that coder was used."""
# Our encoding prepends an 'x:' prefix.
diff --git a/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py b/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py
index 781ec9a8682..51cdc80373d 100644
--- a/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py
+++ b/sdks/python/apache_beam/examples/cookbook/multiple_output_pardo.py
@@ -135,7 +135,9 @@ class CountWords(beam.PTransform):
A PTransform that converts a PCollection containing words into a PCollection
of "word: count" strings.
"""
+
def expand(self, pcoll):
+
def count_ones(word_ones):
(word, ones) = word_ones
return (word, sum(ones))
diff --git a/sdks/python/apache_beam/examples/dataframe/taxiride_it_test.py b/sdks/python/apache_beam/examples/dataframe/taxiride_it_test.py
index f81b7d8cfa6..7ab66e7fd0f 100644
--- a/sdks/python/apache_beam/examples/dataframe/taxiride_it_test.py
+++ b/sdks/python/apache_beam/examples/dataframe/taxiride_it_test.py
@@ -34,6 +34,7 @@ from apache_beam.testing.test_pipeline import TestPipeline
class TaxirideIT(unittest.TestCase):
+
def setUp(self):
self.test_pipeline = TestPipeline(is_integration_test=True)
self.outdir = (
diff --git a/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/main.py b/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/main.py
index ae59b94a1a6..75a88ef23c9 100644
--- a/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/main.py
+++ b/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/main.py
@@ -77,6 +77,7 @@ class PytorchNoBatchModelHandler(PytorchModelHandlerKeyedTensor):
Restricting max_batch_size to 1 means there is only 1 example per `batch`
in the run_inference() call.
"""
+
def batch_elements_kwargs(self):
return {"max_batch_size": 1}
diff --git a/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/pipeline/transformations.py b/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/pipeline/transformations.py
index bb90c5a93cb..dcdf2d3c78e 100644
--- a/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/pipeline/transformations.py
+++ b/sdks/python/apache_beam/examples/inference/anomaly_detection/anomaly_detection_pipeline/pipeline/transformations.py
@@ -60,6 +60,7 @@ def tokenize_sentence(input_dict):
class ModelWrapper(DistilBertModel):
"""Wrapper to DistilBertModel to get embeddings when calling
forward function."""
+
def forward(self, **kwargs):
output = super().forward(**kwargs)
sentence_embedding = (
@@ -123,6 +124,7 @@ class CustomSklearnModelHandlerNumpy(SklearnModelHandlerNumpy):
class NormalizeEmbedding(beam.DoFn):
"""A DoFn for normalization of text embedding."""
+
def process(self, element, *args, **kwargs):
"""
For each element in the input PCollection, normalize the embedding vector, and
@@ -139,6 +141,7 @@ class NormalizeEmbedding(beam.DoFn):
class DecodePubSubMessage(beam.DoFn):
"""A DoFn for decoding PubSub message into a dictionary."""
+
def process(self, element, *args, **kwargs):
"""
For each element in the input PCollection, retrieve the id and decode the bytes into string
@@ -154,6 +157,7 @@ class DecodePubSubMessage(beam.DoFn):
class DecodePrediction(beam.DoFn):
"""A DoFn for decoding the prediction from RunInference."""
+
def process(self, element):
"""
The `process` function takes the output of RunInference and returns a dictionary
@@ -170,6 +174,7 @@ class DecodePrediction(beam.DoFn):
class TriggerEmailAlert(beam.DoFn):
"""A DoFn for sending email using yagmail."""
+
def setup(self):
"""
Opens the cred.json file and initializes the yag SMTP client.
diff --git a/sdks/python/apache_beam/examples/inference/anomaly_detection/write_data_to_pubsub_pipeline/pipeline/utils.py b/sdks/python/apache_beam/examples/inference/anomaly_detection/write_data_to_pubsub_pipeline/pipeline/utils.py
index 66568084349..30b8062e128 100644
--- a/sdks/python/apache_beam/examples/inference/anomaly_detection/write_data_to_pubsub_pipeline/pipeline/utils.py
+++ b/sdks/python/apache_beam/examples/inference/anomaly_detection/write_data_to_pubsub_pipeline/pipeline/utils.py
@@ -55,6 +55,7 @@ def get_dataset(categories: list, split: str = "train"):
class AssignUniqueID(beam.DoFn):
"""A DoFn for assigning Unique ID to each text."""
+
def process(self, element, *args, **kwargs):
uid = str(uuid.uuid4())
yield {"id": uid, "text": element}
@@ -62,6 +63,7 @@ class AssignUniqueID(beam.DoFn):
class ConvertToPubSubMessage(beam.DoFn):
"""A DoFn for converting into PubSub message format."""
+
def process(self, element, *args, **kwargs):
yield PubsubMessage(
data=element["text"].encode("utf-8"), attributes={"id": element["id"]})
diff --git a/sdks/python/apache_beam/examples/inference/huggingface_language_modeling.py b/sdks/python/apache_beam/examples/inference/huggingface_language_modeling.py
index 69c2eacc593..ede4c8753aa 100644
--- a/sdks/python/apache_beam/examples/inference/huggingface_language_modeling.py
+++ b/sdks/python/apache_beam/examples/inference/huggingface_language_modeling.py
@@ -75,6 +75,7 @@ class PostProcessor(beam.DoFn):
The logits are the output of the Model. We can get the word with the highest
probability of being a candidate replacement word by taking the argmax.
"""
+
def __init__(self, tokenizer: AutoTokenizer):
super().__init__()
self.tokenizer = tokenizer
diff --git a/sdks/python/apache_beam/examples/inference/huggingface_question_answering.py b/sdks/python/apache_beam/examples/inference/huggingface_question_answering.py
index 7d4899cc38d..4ac5813a650 100644
--- a/sdks/python/apache_beam/examples/inference/huggingface_question_answering.py
+++ b/sdks/python/apache_beam/examples/inference/huggingface_question_answering.py
@@ -48,6 +48,7 @@ class PostProcessor(beam.DoFn):
Hugging Face Pipeline for Question Answering returns a dictionary
with score, start and end index of answer and the answer.
"""
+
def process(self, result: tuple[str, PredictionResult]) -> Iterable[str]:
text, prediction = result
predicted_answer = prediction.inference['answer']
diff --git a/sdks/python/apache_beam/examples/inference/large_language_modeling/main.py b/sdks/python/apache_beam/examples/inference/large_language_modeling/main.py
index 6fd06ae4d1d..a3adad56ed6 100644
--- a/sdks/python/apache_beam/examples/inference/large_language_modeling/main.py
+++ b/sdks/python/apache_beam/examples/inference/large_language_modeling/main.py
@@ -37,6 +37,7 @@ from transformers import T5ForConditionalGeneration
class Preprocess(beam.DoFn):
+
def __init__(self, tokenizer: AutoTokenizer):
self._tokenizer = tokenizer
@@ -59,6 +60,7 @@ class Preprocess(beam.DoFn):
class Postprocess(beam.DoFn):
+
def __init__(self, tokenizer: AutoTokenizer):
self._tokenizer = tokenizer
diff --git a/sdks/python/apache_beam/examples/inference/milk_quality_prediction_windowing.py b/sdks/python/apache_beam/examples/inference/milk_quality_prediction_windowing.py
index dfec4640fda..71761a49ccf 100644
--- a/sdks/python/apache_beam/examples/inference/milk_quality_prediction_windowing.py
+++ b/sdks/python/apache_beam/examples/inference/milk_quality_prediction_windowing.py
@@ -138,6 +138,7 @@ class MilkQualityAggregation(NamedTuple):
class AggregateMilkQualityResults(beam.CombineFn):
"""Simple aggregation to keep track of the number
of samples with good, bad and medium quality milk."""
+
def create_accumulator(self):
return MilkQualityAggregation(0, 0, 0)
diff --git a/sdks/python/apache_beam/examples/inference/multi_language_inference/multi_language_custom_transform/multi_language_custom_transform/composite_transform.py b/sdks/python/apache_beam/examples/inference/multi_language_inference/multi_language_custom_transform/multi_language_custom_transform/composite_transform.py
index 3e1794bc582..a386acb7708 100644
--- a/sdks/python/apache_beam/examples/inference/multi_language_inference/multi_language_custom_transform/multi_language_custom_transform/composite_transform.py
+++ b/sdks/python/apache_beam/examples/inference/multi_language_inference/multi_language_custom_transform/multi_language_custom_transform/composite_transform.py
@@ -35,6 +35,7 @@ _LOGGER = logging.getLogger(__name__)
class InferenceTransform(ptransform.PTransform):
+
class PytorchModelHandlerKeyedTensorWrapper(PytorchModelHandlerKeyedTensor):
"""Wrapper to PytorchModelHandler to limit batch size to 1.
The tokenized strings generated from BertTokenizer may have different
@@ -43,10 +44,12 @@ class InferenceTransform(ptransform.PTransform):
Restricting max_batch_size to 1 means there is only 1 example per
`batch` in the run_inference() call.
"""
+
def batch_elements_kwargs(self):
return {'max_batch_size': 1}
class Preprocess(beam.DoFn):
+
def __init__(self, tokenizer):
self._tokenizer = tokenizer
logging.info('Starting Preprocess.')
@@ -80,6 +83,7 @@ class InferenceTransform(ptransform.PTransform):
return [(text, tokens)]
class Postprocess(beam.DoFn):
+
def __init__(self, bert_tokenizer):
self.bert_tokenizer = bert_tokenizer
logging.info('Starting Postprocess')
diff --git a/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/main.py b/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/main.py
index 350eec7ed87..4d9187254c1 100644
--- a/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/main.py
+++ b/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/main.py
@@ -74,6 +74,7 @@ class PytorchNoBatchModelHandler(PytorchModelHandlerKeyedTensor):
Restricting max_batch_size to 1 means there is only 1 example per `batch`
in the run_inference() call.
"""
+
def batch_elements_kwargs(self):
return {"max_batch_size": 1}
diff --git a/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/pipeline/transformations.py b/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/pipeline/transformations.py
index 26010516f0c..2e6b2a8e439 100644
--- a/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/pipeline/transformations.py
+++ b/sdks/python/apache_beam/examples/inference/online_clustering/clustering_pipeline/pipeline/transformations.py
@@ -54,6 +54,7 @@ def tokenize_sentence(input_dict):
class ModelWrapper(DistilBertModel):
"""Wrapper to DistilBertModel to get embeddings when calling
forward function."""
+
def forward(self, **kwargs):
output = super().forward(**kwargs)
sentence_embedding = (
@@ -84,6 +85,7 @@ class ModelWrapper(DistilBertModel):
class NormalizeEmbedding(beam.DoFn):
"""A DoFn for normalization of text embedding."""
+
def process(self, element, *args, **kwargs):
"""
For each element in the input PCollection, normalize the embedding vector, and
@@ -99,6 +101,7 @@ class NormalizeEmbedding(beam.DoFn):
class Decode(beam.DoFn):
"""A DoFn for decoding PubSub message into a dictionary."""
+
def process(self, element, *args, **kwargs):
"""
For each element in the input PCollection, retrieve the id and decode the bytes into string
@@ -181,6 +184,7 @@ class StatefulOnlineClustering(beam.DoFn):
class GetUpdates(beam.DoFn):
"""A DoFn for printing the clusters and items belonging to each cluster."""
+
def process(self, element, *args, **kwargs):
"""
Prints and returns clusters with items contained in it
diff --git a/sdks/python/apache_beam/examples/inference/online_clustering/write_data_to_pubsub_pipeline/pipeline/utils.py b/sdks/python/apache_beam/examples/inference/online_clustering/write_data_to_pubsub_pipeline/pipeline/utils.py
index c0007de63d2..d2de237d41c 100644
--- a/sdks/python/apache_beam/examples/inference/online_clustering/write_data_to_pubsub_pipeline/pipeline/utils.py
+++ b/sdks/python/apache_beam/examples/inference/online_clustering/write_data_to_pubsub_pipeline/pipeline/utils.py
@@ -55,6 +55,7 @@ def get_dataset(categories: list, split: str = "train"):
class AssignUniqueID(beam.DoFn):
"""A DoFn for assigning Unique ID to each text."""
+
def process(self, element, *args, **kwargs):
uid = str(uuid.uuid4())
yield {"id": uid, "text": element}
@@ -62,6 +63,7 @@ class AssignUniqueID(beam.DoFn):
class ConvertToPubSubMessage(beam.DoFn):
"""A DoFn for converting into PubSub message format."""
+
def process(self, element, *args, **kwargs):
yield PubsubMessage(
data=element["text"].encode("utf-8"), attributes={"id": element["id"]})
diff --git a/sdks/python/apache_beam/examples/inference/onnx_sentiment_classification.py b/sdks/python/apache_beam/examples/inference/onnx_sentiment_classification.py
index 0e62ab86543..4e9803c890c 100644
--- a/sdks/python/apache_beam/examples/inference/onnx_sentiment_classification.py
+++ b/sdks/python/apache_beam/examples/inference/onnx_sentiment_classification.py
@@ -62,6 +62,7 @@ def filter_empty_lines(text: str) -> Iterator[str]:
class PostProcessor(beam.DoFn):
+
def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]:
filename, prediction_result = element
prediction = np.argmax(prediction_result.inference, axis=0)
@@ -112,6 +113,7 @@ def run(
Restricting max_batch_size to 1 means there is only 1 example per `batch`
in the run_inference() call.
"""
+
def batch_elements_kwargs(self):
return {'max_batch_size': 1}
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 787341263fd..878b5b9763d 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,6 +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]:
filename, prediction_result = element
prediction = torch.argmax(prediction_result.inference, dim=0)
diff --git a/sdks/python/apache_beam/examples/inference/pytorch_image_segmentation.py b/sdks/python/apache_beam/examples/inference/pytorch_image_segmentation.py
index 5e5f77a679c..76f2311ba4f 100644
--- a/sdks/python/apache_beam/examples/inference/pytorch_image_segmentation.py
+++ b/sdks/python/apache_beam/examples/inference/pytorch_image_segmentation.py
@@ -160,6 +160,7 @@ def filter_empty_lines(text: str) -> Iterator[str]:
class PostProcessor(beam.DoFn):
+
def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]:
filename, prediction_result = element
prediction_labels = prediction_result.inference['labels']
diff --git a/sdks/python/apache_beam/examples/inference/pytorch_language_modeling.py b/sdks/python/apache_beam/examples/inference/pytorch_language_modeling.py
index a616998d2c7..a3ef1db7609 100644
--- a/sdks/python/apache_beam/examples/inference/pytorch_language_modeling.py
+++ b/sdks/python/apache_beam/examples/inference/pytorch_language_modeling.py
@@ -78,6 +78,7 @@ class PostProcessor(beam.DoFn):
of the words in BERT’s vocabulary. We can get the word with the highest
probability of being a candidate replacement word by taking the argmax.
"""
+
def __init__(self, bert_tokenizer: BertTokenizer):
super().__init__()
self.bert_tokenizer = bert_tokenizer
@@ -165,6 +166,7 @@ def run(
Restricting max_batch_size to 1 means there is only 1 example per `batch`
in the run_inference() call.
"""
+
def batch_elements_kwargs(self):
return {'max_batch_size': 1}
@@ -181,18 +183,19 @@ def run(
bert_tokenizer = BertTokenizer.from_pretrained(known_args.bert_tokenizer)
if not known_args.input:
- text = (pipeline | 'CreateSentences' >> beam.Create([
- 'The capital of France is Paris .',
- 'It is raining cats and dogs .',
- 'He looked up and saw the sun and stars .',
- 'Today is Monday and tomorrow is Tuesday .',
- 'There are 5 coconuts on this palm tree .',
- 'The richest person in the world is not here .',
- 'Malls are amazing places to shop because you can find everything you need under one roof .', # pylint: disable=line-too-long
- 'This audiobook is sure to liquefy your brain .',
- 'The secret ingredient to his wonderful life was gratitude .',
- 'The biggest animal in the world is the whale .',
- ]))
+ text = (
+ pipeline | 'CreateSentences' >> beam.Create([
+ 'The capital of France is Paris .',
+ 'It is raining cats and dogs .',
+ 'He looked up and saw the sun and stars .',
+ 'Today is Monday and tomorrow is Tuesday .',
+ 'There are 5 coconuts on this palm tree .',
+ 'The richest person in the world is not here .',
+ 'Malls are amazing places to shop because you can find everything you need under one roof .', # pylint: disable=line-too-long
+ 'This audiobook is sure to liquefy your brain .',
+ 'The secret ingredient to his wonderful life was gratitude .',
+ 'The biggest animal in the world is the whale .',
+ ]))
else:
text = (
pipeline | 'ReadSentences' >> beam.io.ReadFromText(known_args.input))
diff --git a/sdks/python/apache_beam/examples/inference/pytorch_model_per_key_image_segmentation.py b/sdks/python/apache_beam/examples/inference/pytorch_model_per_key_image_segmentation.py
index 18c4c3e653b..0bc5ed266df 100644
--- a/sdks/python/apache_beam/examples/inference/pytorch_model_per_key_image_segmentation.py
+++ b/sdks/python/apache_beam/examples/inference/pytorch_model_per_key_image_segmentation.py
@@ -166,6 +166,7 @@ def filter_empty_lines(text: str) -> Iterator[str]:
class KeyExamplesForEachModelType(beam.DoFn):
"""Duplicate data to run against each model type"""
+
def process(
self, element: tuple[torch.Tensor,
str]) -> Iterable[tuple[str, torch.Tensor]]:
@@ -174,6 +175,7 @@ class KeyExamplesForEachModelType(beam.DoFn):
class PostProcessor(beam.DoFn):
+
def process(
self, element: tuple[str, PredictionResult]) -> tuple[torch.Tensor, str]:
model, prediction_result = element
@@ -183,6 +185,7 @@ class PostProcessor(beam.DoFn):
class FormatResults(beam.DoFn):
+
def process(self, element):
_, filename_prediction = element
predictions = filename_prediction['predictions']
diff --git a/sdks/python/apache_beam/examples/inference/run_inference_side_inputs.py b/sdks/python/apache_beam/examples/inference/run_inference_side_inputs.py
index 755eff17c16..42cd7596bcc 100644
--- a/sdks/python/apache_beam/examples/inference/run_inference_side_inputs.py
+++ b/sdks/python/apache_beam/examples/inference/run_inference_side_inputs.py
@@ -38,22 +38,26 @@ from apache_beam.transforms.userstate import CombiningValueStateSpec
# create some fake models which returns different inference results.
class FakeModelDefault:
+
def predict(self, example: int) -> int:
return example
class FakeModelAdd(FakeModelDefault):
+
def predict(self, example: int) -> int:
return example + 1
class FakeModelSub(FakeModelDefault):
+
def predict(self, example: int) -> int:
return example - 1
class FakeModelHandlerReturnsPredictionResult(
base.ModelHandler[int, base.PredictionResult, FakeModelDefault]):
+
def __init__(self, clock=None, model_id='model_default'):
self.model_id = model_id
self._fake_clock = clock
@@ -96,6 +100,7 @@ def run(argv=None, save_main_session=True):
options.view_as(SetupOptions).save_main_session = save_main_session
class GetModel(beam.DoFn):
+
def process(self, element) -> Iterable[base.ModelMetadata]:
if time.time() > mid_ts:
yield base.ModelMetadata(
diff --git a/sdks/python/apache_beam/examples/inference/runinference_metrics/pipeline/transformations.py b/sdks/python/apache_beam/examples/inference/runinference_metrics/pipeline/transformations.py
index e7f6f9d4468..ccbf6753078 100644
--- a/sdks/python/apache_beam/examples/inference/runinference_metrics/pipeline/transformations.py
+++ b/sdks/python/apache_beam/examples/inference/runinference_metrics/pipeline/transformations.py
@@ -27,6 +27,7 @@ from transformers import DistilBertTokenizer
class CustomPytorchModelHandlerKeyedTensor(PytorchModelHandlerKeyedTensor):
"""Wrapper around PytorchModelHandlerKeyedTensor to load a model on CPU."""
+
def load_model(self) -> torch.nn.Module:
"""Loads and initializes a Pytorch model for processing."""
model = self._model_class(**self._model_params)
@@ -43,6 +44,7 @@ class HuggingFaceStripBatchingWrapper(DistilBertForSequenceClassification):
as a list of dicts instead of a dict of lists. Another workaround
can be found here where they disable batching instead.
https://github.com/apache/beam/blob/master/sdks/python/apache_beam/examples/inference/pytorch_language_modeling.py"""
+
def forward(self, **kwargs):
output = super().forward(**kwargs)
return [dict(zip(output, v)) for v in zip(*output.values())]
@@ -50,6 +52,7 @@ class HuggingFaceStripBatchingWrapper(DistilBertForSequenceClassification):
class Tokenize(beam.DoFn):
"""A DoFn for tokenizing texts"""
+
def __init__(self, model_name: str):
"""Initialises a tokenizer based on the model_name"""
self._model_name = model_name
@@ -75,6 +78,7 @@ class Tokenize(beam.DoFn):
class PostProcessor(beam.DoFn):
"""Postprocess the RunInference output"""
+
def process(self, element):
"""
Takes the input text and the prediction result, and returns a dictionary
diff --git a/sdks/python/apache_beam/examples/inference/sklearn_japanese_housing_regression.py b/sdks/python/apache_beam/examples/inference/sklearn_japanese_housing_regression.py
index 0a527e88dec..8b85af3a05b 100644
--- a/sdks/python/apache_beam/examples/inference/sklearn_japanese_housing_regression.py
+++ b/sdks/python/apache_beam/examples/inference/sklearn_japanese_housing_regression.py
@@ -88,6 +88,7 @@ def sort_by_features(dataframe, max_size):
class LoadDataframe(beam.DoFn):
+
def process(self, file_name: str) -> Iterable[pandas.DataFrame]:
""" Loads data files as a pandas dataframe."""
file = FileSystems.open(file_name, 'rb')
diff --git a/sdks/python/apache_beam/examples/inference/sklearn_mnist_classification.py b/sdks/python/apache_beam/examples/inference/sklearn_mnist_classification.py
index d7d08e294e9..6f3890c3c0a 100644
--- a/sdks/python/apache_beam/examples/inference/sklearn_mnist_classification.py
+++ b/sdks/python/apache_beam/examples/inference/sklearn_mnist_classification.py
@@ -51,6 +51,7 @@ class PostProcessor(beam.DoFn):
"""Process the PredictionResult to get the predicted label.
Returns a comma separated string with true label and predicted label.
"""
+
def process(self, element: tuple[int, PredictionResult]) -> Iterable[str]:
label, prediction_result = element
prediction = prediction_result.inference
diff --git a/sdks/python/apache_beam/examples/inference/tensorflow_imagenet_segmentation.py b/sdks/python/apache_beam/examples/inference/tensorflow_imagenet_segmentation.py
index b44d775f4ad..a24d8c6b3db 100644
--- a/sdks/python/apache_beam/examples/inference/tensorflow_imagenet_segmentation.py
+++ b/sdks/python/apache_beam/examples/inference/tensorflow_imagenet_segmentation.py
@@ -37,6 +37,7 @@ class PostProcessor(beam.DoFn):
"""Process the PredictionResult to get the predicted label.
Returns predicted label.
"""
+
def setup(self):
labels_path = tf.keras.utils.get_file(
'ImageNetLabels.txt',
diff --git a/sdks/python/apache_beam/examples/inference/tensorflow_mnist_classification.py b/sdks/python/apache_beam/examples/inference/tensorflow_mnist_classification.py
index bf85bb1aef1..f5f9b4fa8c9 100644
--- a/sdks/python/apache_beam/examples/inference/tensorflow_mnist_classification.py
+++ b/sdks/python/apache_beam/examples/inference/tensorflow_mnist_classification.py
@@ -45,6 +45,7 @@ class PostProcessor(beam.DoFn):
"""Process the PredictionResult to get the predicted label.
Returns a comma separated string with true label and predicted label.
"""
+
def process(self, element: tuple[int, PredictionResult]) -> Iterable[str]:
label, prediction_result = element
prediction = numpy.argmax(prediction_result.inference, axis=0)
diff --git a/sdks/python/apache_beam/examples/inference/tensorrt_object_detection.py b/sdks/python/apache_beam/examples/inference/tensorrt_object_detection.py
index 677d36b9b76..39e1d625375 100644
--- a/sdks/python/apache_beam/examples/inference/tensorrt_object_detection.py
+++ b/sdks/python/apache_beam/examples/inference/tensorrt_object_detection.py
@@ -167,6 +167,7 @@ class PostProcessor(beam.DoFn):
an integer that we can transform into actual string class using
COCO_OBJ_DET_CLASSES as reference.
"""
+
def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]:
key, prediction_result = element
filename, im_width, im_height = key
diff --git a/sdks/python/apache_beam/examples/inference/tensorrt_text_classification.py b/sdks/python/apache_beam/examples/inference/tensorrt_text_classification.py
index a5cda68fd79..e101e7d2de8 100644
--- a/sdks/python/apache_beam/examples/inference/tensorrt_text_classification.py
+++ b/sdks/python/apache_beam/examples/inference/tensorrt_text_classification.py
@@ -43,6 +43,7 @@ class Preprocess(beam.DoFn):
The input sentences are tokenized because the
model is expecting tokens.
"""
+
def __init__(self, tokenizer: AutoTokenizer):
self._tokenizer = tokenizer
@@ -59,6 +60,7 @@ class Postprocess(beam.DoFn):
We can get the class label by getting the index of
maximum logit using argmax.
"""
+
def __init__(self, tokenizer: AutoTokenizer):
self._tokenizer = tokenizer
diff --git a/sdks/python/apache_beam/examples/inference/tfx_bsl/build_tensorflow_model.py b/sdks/python/apache_beam/examples/inference/tfx_bsl/build_tensorflow_model.py
index 9230f84955e..31fd8c922f5 100644
--- a/sdks/python/apache_beam/examples/inference/tfx_bsl/build_tensorflow_model.py
+++ b/sdks/python/apache_beam/examples/inference/tfx_bsl/build_tensorflow_model.py
@@ -54,6 +54,7 @@ class TFModelWrapperWithSignature(tf.keras.Model):
saved_model_spec=saved_model_spec)
model_handler = CreateModelHandler(inferece_spec_type)
"""
+
def __init__(
self,
model,
diff --git a/sdks/python/apache_beam/examples/inference/tfx_bsl/tensorflow_image_classification.py b/sdks/python/apache_beam/examples/inference/tfx_bsl/tensorflow_image_classification.py
index 5df0b51e36d..ed78534dd43 100644
--- a/sdks/python/apache_beam/examples/inference/tfx_bsl/tensorflow_image_classification.py
+++ b/sdks/python/apache_beam/examples/inference/tfx_bsl/tensorflow_image_classification.py
@@ -95,6 +95,7 @@ def convert_image_to_example_proto(tensor: tf.Tensor) -> tf.train.Example:
class ProcessInferenceToString(beam.DoFn):
+
def process(
self, element: tuple[str,
prediction_log_pb2.PredictionLog]) -> Iterable[str]:
diff --git a/sdks/python/apache_beam/examples/inference/tfx_bsl/tfx_bsl_inference_it_test.py b/sdks/python/apache_beam/examples/inference/tfx_bsl/tfx_bsl_inference_it_test.py
index d72794df4f7..4b437677585 100644
--- a/sdks/python/apache_beam/examples/inference/tfx_bsl/tfx_bsl_inference_it_test.py
+++ b/sdks/python/apache_beam/examples/inference/tfx_bsl/tfx_bsl_inference_it_test.py
@@ -60,6 +60,7 @@ def process_outputs(filepath):
tfx_bsl is None, 'Missing dependencies. '
'Test depends on tfx_bsl')
class TFXRunInferenceTests(unittest.TestCase):
+
@pytest.mark.uses_tensorflow
@pytest.mark.it_postcommit
def test_tfx_run_inference_mobilenetv2(self):
diff --git a/sdks/python/apache_beam/examples/inference/vertex_ai_image_classification.py b/sdks/python/apache_beam/examples/inference/vertex_ai_image_classification.py
index 20312e7d3c8..55c28b24447 100644
--- a/sdks/python/apache_beam/examples/inference/vertex_ai_image_classification.py
+++ b/sdks/python/apache_beam/examples/inference/vertex_ai_image_classification.py
@@ -117,6 +117,7 @@ def preprocess_image(data: bytes) -> list[float]:
class PostProcessor(beam.DoFn):
+
def process(self, element: tuple[str, PredictionResult]) -> Iterable[str]:
img_name, prediction_result = element
prediction_vals = prediction_result.inference
diff --git a/sdks/python/apache_beam/examples/inference/vllm_text_completion.py b/sdks/python/apache_beam/examples/inference/vllm_text_completion.py
index 2708c0f3d1a..f1f0b769aff 100644
--- a/sdks/python/apache_beam/examples/inference/vllm_text_completion.py
+++ b/sdks/python/apache_beam/examples/inference/vllm_text_completion.py
@@ -116,6 +116,7 @@ def parse_known_args(argv):
class PostProcessor(beam.DoFn):
+
def process(self, element: PredictionResult) -> Iterable[str]:
yield str(element.example) + ": " + str(element.inference)
diff --git a/sdks/python/apache_beam/examples/inference/xgboost_iris_classification.py b/sdks/python/apache_beam/examples/inference/xgboost_iris_classification.py
index 498511a5a2c..660a745f7e6 100644
--- a/sdks/python/apache_beam/examples/inference/xgboost_iris_classification.py
+++ b/sdks/python/apache_beam/examples/inference/xgboost_iris_classification.py
@@ -46,6 +46,7 @@ class PostProcessor(beam.DoFn):
"""Process the PredictionResult to get the predicted label.
Returns a comma separated string with true label and predicted label.
"""
+
def process(self, element: tuple[int, PredictionResult]) -> Iterable[str]:
label, prediction_result = element
prediction = prediction_result.inference
@@ -101,8 +102,8 @@ def load_sklearn_iris_test_data(
dataset['data'], dataset['target'], test_size=.2, random_state=seed)
if split:
- return [(index, data_type(sample.reshape(1, -1))) for index,
- sample in enumerate(x_test)]
+ return [(index, data_type(sample.reshape(1, -1)))
+ for index, sample in enumerate(x_test)]
return [(0, data_type(x_test))]
diff --git a/sdks/python/apache_beam/examples/ml_transform/ml_transform_it_test.py b/sdks/python/apache_beam/examples/ml_transform/ml_transform_it_test.py
index 96fb3f77567..e28dc3313fa 100644
--- a/sdks/python/apache_beam/examples/ml_transform/ml_transform_it_test.py
+++ b/sdks/python/apache_beam/examples/ml_transform/ml_transform_it_test.py
@@ -59,6 +59,7 @@ def _publish_metrics(pipeline, metric_value, metrics_table, metric_name):
@pytest.mark.uses_tft
class LargeMovieReviewDatasetProcessTest(unittest.TestCase):
+
def test_process_large_movie_review_dataset(self):
input_data_dir = 'gs://apache-beam-ml/datasets/aclImdb'
artifact_location = os.path.join(_OUTPUT_GCS_BUCKET_ROOT, uuid.uuid4().hex)
diff --git a/sdks/python/apache_beam/examples/ml_transform/vocab_tfidf_processing.py b/sdks/python/apache_beam/examples/ml_transform/vocab_tfidf_processing.py
index b8ae61ce51e..f0d5d1c652b 100644
--- a/sdks/python/apache_beam/examples/ml_transform/vocab_tfidf_processing.py
+++ b/sdks/python/apache_beam/examples/ml_transform/vocab_tfidf_processing.py
@@ -66,6 +66,7 @@ def Shuffle(pcoll):
class ReadAndShuffleData(beam.PTransform):
+
def __init__(self, pos_file_pattern, neg_file_pattern):
self.pos_file_pattern = pos_file_pattern
self.neg_file_pattern = neg_file_pattern
@@ -94,8 +95,7 @@ class ReadAndShuffleData(beam.PTransform):
shuffled_examples
| beam.Map(
lambda label_review: {
- REVIEW_COLUMN: label_review[0],
- LABEL_COLUMN: label_review[1],
+ REVIEW_COLUMN: label_review[0], LABEL_COLUMN: label_review[1],
RAW_DATA_KEY: label_review[0]
}))
@@ -144,6 +144,7 @@ def preprocess_data(
class MapTFIDFScoreToVocab(beam.DoFn):
+
def __init__(self, artifact_location):
self.artifact_location = artifact_location
diff --git a/sdks/python/apache_beam/examples/online_clustering.py b/sdks/python/apache_beam/examples/online_clustering.py
index 9746995dce7..5c71bc52353 100644
--- a/sdks/python/apache_beam/examples/online_clustering.py
+++ b/sdks/python/apache_beam/examples/online_clustering.py
@@ -37,6 +37,7 @@ from apache_beam.transforms.userstate import ReadModifyWriteStateSpec
class SaveModel(core.DoFn):
"""Saves trained clustering model to persistent storage"""
+
def __init__(self, checkpoints_path: str):
self.checkpoints_path = checkpoints_path
@@ -64,6 +65,7 @@ class SaveModel(core.DoFn):
class AssignClusterLabelsFn(core.DoFn):
"""Takes a trained model and input data and labels
all data instances using the trained model."""
+
def process(self, batch, model, model_id):
cluster_labels = model.predict(batch)
for e, i in zip(batch, cluster_labels):
@@ -72,6 +74,7 @@ class AssignClusterLabelsFn(core.DoFn):
class SelectLatestModelState(core.CombineFn):
"""Selects that latest version of a model after training"""
+
def create_accumulator(self):
# create and initialise accumulator
return None, 0
@@ -178,6 +181,7 @@ class OnlineKMeans(ClusteringAlgorithm):
class ConvertToNumpyArray(core.DoFn):
"""Helper function to convert incoming data
to numpy arrays that are accepted by sklearn"""
+
def process(self, element, *args, **kwargs):
if isinstance(element, (tuple, list)):
yield np.array(element)
@@ -190,6 +194,7 @@ class ConvertToNumpyArray(core.DoFn):
class ClusteringPreprocessing(ptransform.PTransform):
+
def __init__(
self, n_clusters: int, batch_size: int, is_batched: bool = False):
""" Preprocessing for Clustering Transformation
@@ -235,6 +240,7 @@ class ClusteringPreprocessing(ptransform.PTransform):
class OnlineClustering(ptransform.PTransform):
+
def __init__(
self,
clustering_algorithm,
@@ -303,6 +309,7 @@ class OnlineClustering(ptransform.PTransform):
class AssignClusterLabelsRunInference(ptransform.PTransform):
+
def __init__(self, checkpoints_path):
super().__init__()
self.clustering_model = SklearnModelHandlerNumpy(
@@ -318,6 +325,7 @@ class AssignClusterLabelsRunInference(ptransform.PTransform):
class AssignClusterLabelsInMemoryModel(ptransform.PTransform):
+
def __init__(
self, model, n_clusters, batch_size, is_batched=False, model_id=None):
self.model = model
diff --git a/sdks/python/apache_beam/examples/per_entity_training.py b/sdks/python/apache_beam/examples/per_entity_training.py
index e2b267492b2..7eef6d47c20 100644
--- a/sdks/python/apache_beam/examples/per_entity_training.py
+++ b/sdks/python/apache_beam/examples/per_entity_training.py
@@ -44,6 +44,7 @@ from apache_beam.options.pipeline_options import SetupOptions
class CreateKey(beam.DoFn):
+
def process(self, element, *args, **kwargs):
# 3rd column of the dataset is Education
idx = 3
@@ -62,6 +63,7 @@ def custom_filter(element):
class PrepareDataforTraining(beam.DoFn):
"""Preprocess data in a format suitable for training."""
+
def process(self, element, *args, **kwargs):
key, values = element
#Convert to dataframe
@@ -82,6 +84,7 @@ class TrainModel(beam.DoFn):
normalizes numerical columns and then
fits a decision tree classifier.
"""
+
def process(self, element, *args, **kwargs):
X, y, cat_ix, num_ix, key = element
steps = [('c', OneHotEncoder(handle_unknown='ignore'), cat_ix),
@@ -95,6 +98,7 @@ class TrainModel(beam.DoFn):
class ModelSink(fileio.FileSink):
+
def open(self, fh):
self._fh = fh
diff --git a/sdks/python/apache_beam/examples/snippets/snippets.py b/sdks/python/apache_beam/examples/snippets/snippets.py
index c849af4a00b..1833ae64386 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets.py
@@ -80,6 +80,7 @@ class SnippetUtils(object):
This is as close as we can get to have code snippets that are
executed and are also ready to presented in webdocs.
"""
+
def __init__(self, renames):
self.renames = renames
@@ -199,6 +200,7 @@ def pipeline_options_remote():
from apache_beam.options.pipeline_options import PipelineOptions
class MyOptions(PipelineOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
parser.add_argument('--input')
@@ -256,6 +258,7 @@ def pipeline_options_local():
from apache_beam.options.pipeline_options import PipelineOptions
class MyOptions(PipelineOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
parser.add_argument(
@@ -332,6 +335,7 @@ def pipeline_logging(lines, output):
import logging
class ExtractWordsFn(beam.DoFn):
+
def process(self, element):
words = re.findall(r'[A-Za-z\']+', element)
for word in words:
@@ -360,12 +364,14 @@ def pipeline_monitoring():
import apache_beam as beam
class ExtractWordsFn(beam.DoFn):
+
def process(self, element):
words = re.findall(r'[A-Za-z\']+', element)
for word in words:
yield word
class FormatCountsFn(beam.DoFn):
+
def process(self, element):
word, count = element
yield '%s: %s' % (word, count)
@@ -419,6 +425,7 @@ def examples_wordcount_templated():
# [START example_wordcount_templated]
class WordcountTemplatedOptions(PipelineOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
# Use add_value_provider_argument for arguments to be templatable
@@ -521,11 +528,13 @@ def examples_ptransforms_templated(renames):
from apache_beam.options.value_provider import StaticValueProvider
class TemplatedUserOptions(PipelineOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument('--templated_int', type=int)
class MySumFn(beam.DoFn):
+
def __init__(self, templated_int):
self.templated_int = templated_int
@@ -556,6 +565,7 @@ def examples_ptransforms_templated(renames):
# Defining a new source.
# [START model_custom_source_new_source]
class CountingSource(iobase.BoundedSource):
+
def __init__(self, count):
self.records_read = Metrics.counter(self.__class__, 'recordsRead')
self._count = count
@@ -609,6 +619,7 @@ class _CountingSource(CountingSource):
# [START model_custom_source_new_ptransform]
class ReadFromCountingSource(PTransform):
+
def __init__(self, count):
super().__init__()
self._count = count
@@ -682,6 +693,7 @@ def model_custom_source(count):
#
# [START model_custom_sink_new_sink]
class SimpleKVSink(iobase.Sink):
+
def __init__(self, simplekv, url, final_table_name):
self._simplekv = simplekv
self._url = url
@@ -710,6 +722,7 @@ class SimpleKVSink(iobase.Sink):
# Defining a writer for the new sink.
# [START model_custom_sink_new_writer]
class SimpleKVWriter(iobase.Writer):
+
def __init__(self, simplekv, access_token, table_name):
self._simplekv = simplekv
self._access_token = access_token
@@ -730,6 +743,7 @@ class SimpleKVWriter(iobase.Writer):
# [START model_custom_sink_new_ptransform]
class WriteToKVSink(PTransform):
+
def __init__(self, simplekv, url, final_table_name):
self._simplekv = simplekv
super().__init__()
@@ -810,6 +824,7 @@ def model_custom_sink(
def model_textio(renames):
"""Using a Read and Write transform to read/write text files."""
+
def filter_words(x):
import re
return re.findall(r'[A-Za-z\']+', x)
@@ -1317,6 +1332,7 @@ def model_join_using_side_inputs(
# [START model_library_transforms_keys]
class Keys(beam.PTransform):
+
def expand(self, pcoll):
return pcoll | 'Keys' >> beam.Map(lambda k_v: k_v[0])
@@ -1327,6 +1343,7 @@ class Keys(beam.PTransform):
# [START model_library_transforms_count]
class Count(beam.PTransform):
+
def expand(self, pcoll):
return (
pcoll
@@ -1364,11 +1381,13 @@ def accessing_valueprovider_info_after_run():
from apache_beam.options.value_provider import RuntimeValueProvider
class MyOptions(PipelineOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument('--string_value', type=str)
class LogValueProvidersFn(beam.DoFn):
+
def __init__(self, string_vp):
self.string_vp = string_vp
@@ -1569,6 +1588,7 @@ def sdf_basic_example():
# [START SDF_BasicExample]
class FileToWordsRestrictionProvider(beam.transforms.core.RestrictionProvider
):
+
def initial_restriction(self, file_name):
return OffsetRange(0, os.stat(file_name).st_size)
@@ -1576,6 +1596,7 @@ def sdf_basic_example():
return beam.io.restriction_trackers.OffsetRestrictionTracker()
class FileToWordsFn(beam.DoFn):
+
def process(
self,
file_name,
@@ -1603,6 +1624,7 @@ def sdf_basic_example_with_splitting():
# [START SDF_BasicExampleWithSplitting]
class FileToWordsRestrictionProvider(beam.transforms.core.RestrictionProvider
):
+
def split(self, file_name, restriction):
# Compute and output 64 MiB size ranges to process in parallel
split_size = 64 * (1 << 20)
@@ -1624,6 +1646,7 @@ def sdf_sdk_initiated_checkpointing():
# [START SDF_UserInitiatedCheckpoint]
class MySplittableDoFn(beam.DoFn):
+
def process(
self,
element,
@@ -1657,6 +1680,7 @@ def sdf_get_size():
# The RestrictionProvider is responsible for calculating the size of given
# restriction.
class MyRestrictionProvider(beam.transforms.core.RestrictionProvider):
+
def restriction_size(self, file_name, restriction):
weight = 2 if "expensiveRecords" in file_name else 1
return restriction.size() * weight
@@ -1665,6 +1689,7 @@ def sdf_get_size():
def sdf_bad_try_claim_loop():
+
class FileToWordsRestrictionProvider(object):
pass
@@ -1672,6 +1697,7 @@ def sdf_bad_try_claim_loop():
# [START SDF_BadTryClaimLoop]
class BadTryClaimLoop(beam.DoFn):
+
def process(
self,
file_name,
@@ -1702,12 +1728,14 @@ def sdf_custom_watermark_estimator():
# (Optional) Define a custom watermark state type to save information between
# bundle processing rounds.
class MyCustomerWatermarkEstimatorState(object):
+
def __init__(self, element, restriction):
# Store data necessary for future watermark computations
pass
# Define a WatermarkEstimator
class MyCustomWatermarkEstimator(WatermarkEstimator):
+
def __init__(self, estimator_state):
self.state = estimator_state
@@ -1727,6 +1755,7 @@ def sdf_custom_watermark_estimator():
# Then, a WatermarkEstimatorProvider needs to be created for this
# WatermarkEstimator
class MyWatermarkEstimatorProvider(WatermarkEstimatorProvider):
+
def initial_estimator_state(self, element, restriction):
return MyCustomerWatermarkEstimatorState(element, restriction)
@@ -1735,6 +1764,7 @@ def sdf_custom_watermark_estimator():
# Finally, define the SDF using your estimator.
class MySplittableDoFn(beam.DoFn):
+
def process(
self,
element,
@@ -1750,6 +1780,7 @@ def sdf_custom_watermark_estimator():
def sdf_truncate():
# [START SDF_Truncate]
class MyRestrictionProvider(beam.transforms.core.RestrictionProvider):
+
def truncate(self, file_name, restriction):
if "optional" in file_name:
# Skip optional files
@@ -1764,6 +1795,7 @@ def bundle_finalize():
# [START BundleFinalize]
class MySplittableDoFn(beam.DoFn):
+
def process(self, element, bundle_finalizer=beam.DoFn.BundleFinalizerParam):
# ... produce output ...
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_debugging.py b/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_debugging.py
index 8efa3d827bd..ac3aac5ae0b 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_debugging.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_debugging.py
@@ -63,6 +63,7 @@ def examples_wordcount_debugging(renames):
class FilterTextFn(beam.DoFn):
"""A DoFn that filters for a specific key based on a regular expression."""
+
def __init__(self, pattern):
self.pattern = pattern
# A custom metric can track values in your pipeline as it runs. Create
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_wordcount.py b/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_wordcount.py
index 30e6fe7dced..021094e3e9a 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_wordcount.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_examples_wordcount_wordcount.py
@@ -88,6 +88,7 @@ def examples_wordcount_wordcount():
# [START examples_wordcount_wordcount_dofn]
class FormatAsTextFn(beam.DoFn):
+
def process(self, element):
word, count = element
yield '%s: %s' % (word, count)
diff --git a/sdks/python/apache_beam/examples/snippets/snippets_test.py b/sdks/python/apache_beam/examples/snippets/snippets_test.py
index 54a57673b5f..3c4aba1815a 100644
--- a/sdks/python/apache_beam/examples/snippets/snippets_test.py
+++ b/sdks/python/apache_beam/examples/snippets/snippets_test.py
@@ -92,6 +92,7 @@ except ImportError:
class ParDoTest(unittest.TestCase):
"""Tests for model/par-do."""
+
def test_pardo(self):
# Note: "words" and "ComputeWordLengthFn" are referenced by name in
# the text of the doc.
@@ -100,6 +101,7 @@ class ParDoTest(unittest.TestCase):
# [START model_pardo_pardo]
class ComputeWordLengthFn(beam.DoFn):
+
def process(self, element):
return [len(element)]
@@ -116,6 +118,7 @@ class ParDoTest(unittest.TestCase):
# [START model_pardo_yield]
class ComputeWordLengthFn(beam.DoFn):
+
def process(self, element):
yield len(element)
@@ -209,6 +212,7 @@ class ParDoTest(unittest.TestCase):
# [START model_pardo_side_input_dofn]
class FilterUsingLength(beam.DoFn):
+
def process(self, element, lower_bound, upper_bound=float('inf')):
if lower_bound <= len(element) <= upper_bound:
yield element
@@ -220,6 +224,7 @@ class ParDoTest(unittest.TestCase):
def test_pardo_with_tagged_outputs(self):
# [START model_pardo_emitting_values_on_tagged_outputs]
class ProcessWords(beam.DoFn):
+
def process(self, element, cutoff_length, marker):
if len(element) <= cutoff_length:
# Emit this short word to the main output.
@@ -289,6 +294,7 @@ class ParDoTest(unittest.TestCase):
class TypeHintsTest(unittest.TestCase):
+
def test_bad_types(self):
# [START type_hints_missing_define_numbers]
p = TestPipeline()
@@ -322,6 +328,7 @@ class TypeHintsTest(unittest.TestCase):
# [START type_hints_do_fn]
@beam.typehints.with_input_types(int)
class FilterEvensDoFn(beam.DoFn):
+
def process(self, element):
if element % 2 == 0:
yield element
@@ -340,6 +347,7 @@ class TypeHintsTest(unittest.TestCase):
@beam.typehints.with_input_types(T)
@beam.typehints.with_output_types(Tuple[int, T])
class MyTransform(beam.PTransform):
+
def expand(self, pcoll):
return pcoll | beam.Map(lambda x: (len(x), x))
@@ -362,6 +370,7 @@ class TypeHintsTest(unittest.TestCase):
# pylint: disable=expression-not-assigned
# pylint: disable=unused-variable
class FilterEvensDoFn(beam.DoFn):
+
def process(self, element):
if element % 2 == 0:
yield element
@@ -383,6 +392,7 @@ class TypeHintsTest(unittest.TestCase):
from typing import Iterable
class TypedFilterEvensDoFn(beam.DoFn):
+
def process(self, element: int) -> Iterable[int]:
if element % 2 == 0:
yield element
@@ -397,6 +407,7 @@ class TypeHintsTest(unittest.TestCase):
from typing import List, Optional
class FilterEvensDoubleDoFn(beam.DoFn):
+
def process(self, element: int) -> Optional[List[int]]:
if element % 2 == 0:
return [element, element]
@@ -420,6 +431,7 @@ class TypeHintsTest(unittest.TestCase):
from apache_beam.pvalue import PCollection
class IntToStr(beam.PTransform):
+
def expand(self, pcoll: PCollection[int]) -> PCollection[str]:
return pcoll | beam.Map(lambda elem: str(elem))
@@ -463,11 +475,13 @@ class TypeHintsTest(unittest.TestCase):
from typing import Tuple
class Player(object):
+
def __init__(self, team, name):
self.team = team
self.name = name
class PlayerCoder(beam.coders.Coder):
+
def encode(self, player):
return ('%s:%s' % (player.team, player.name)).encode('utf-8')
@@ -502,11 +516,13 @@ class SnippetsTest(unittest.TestCase):
To be used for testing.
"""
+
def __init__(self, file_to_read=None, compression_type=None):
self.file_to_read = file_to_read
self.compression_type = compression_type
class ReadDoFn(beam.DoFn):
+
def __init__(self, file_to_read, compression_type):
self.file_to_read = file_to_read
self.compression_type = compression_type
@@ -541,10 +557,12 @@ class SnippetsTest(unittest.TestCase):
To be used for testing.
"""
+
def __init__(self, file_to_write=None, file_name_suffix=''):
self.file_to_write = file_to_write
class WriteDoFn(beam.DoFn):
+
def __init__(self, file_to_write):
self.file_to_write = file_to_write
self.file_obj = None
@@ -643,6 +661,7 @@ class SnippetsTest(unittest.TestCase):
tempdir_name = tempfile.mkdtemp()
class SimpleKV(object):
+
def __init__(self, tmp_dir):
self._dummy_token = 'dummy_token'
self._tmp_dir = tmp_dir
@@ -841,6 +860,7 @@ class SnippetsTest(unittest.TestCase):
@mock.patch('apache_beam.io.ReadFromPubSub')
@mock.patch('apache_beam.io.WriteToPubSub')
def test_examples_wordcount_streaming(self, *unused_mocks):
+
def FakeReadFromPubSub(topic=None, subscription=None, values=None):
expected_topic = topic
expected_subscription = subscription
@@ -853,6 +873,7 @@ class SnippetsTest(unittest.TestCase):
return _inner
class AssertTransform(beam.PTransform):
+
def __init__(self, matcher):
self.matcher = matcher
@@ -998,8 +1019,8 @@ class SnippetsTest(unittest.TestCase):
]
# [END model_group_by_key_cogroupbykey_tuple_formatted_outputs]
expected_results = [
- '%s; %s; %s' % (name, info['emails'], info['phones']) for name,
- info in results
+ '%s; %s; %s' % (name, info['emails'], info['phones'])
+ for name, info in results
]
self.assertEqual(expected_results, formatted_results)
self.assertEqual(formatted_results, self.get_output(result_path))
@@ -1017,6 +1038,7 @@ class SnippetsTest(unittest.TestCase):
# [START metrics_usage_example]
class FilterTextFn(beam.DoFn):
"""A DoFn that filters for a specific key based on a regex."""
+
def __init__(self, pattern):
self.pattern = pattern
# A custom metric can track values in your pipeline as it runs. Create
@@ -1192,6 +1214,7 @@ class SnippetsTest(unittest.TestCase):
class CombineTest(unittest.TestCase):
"""Tests for model/combine."""
+
def test_global_sum(self):
pc = [1, 2, 3]
# [START global_sum]
@@ -1259,6 +1282,7 @@ class CombineTest(unittest.TestCase):
# [START combine_custom_average_define]
class AverageFn(beam.CombineFn):
+
def create_accumulator(self):
return (0.0, 0)
@@ -1380,6 +1404,7 @@ class CombineTest(unittest.TestCase):
# [START setting_timestamp]
class AddTimestampDoFn(beam.DoFn):
+
def process(self, element):
# Extract the numeric Unix seconds-since-epoch timestamp to be
# associated with the current log entry.
@@ -1403,10 +1428,12 @@ class CombineTest(unittest.TestCase):
class PTransformTest(unittest.TestCase):
"""Tests for PTransform."""
+
def test_composite(self):
# [START model_composite_transform]
class ComputeWordLengths(beam.PTransform):
+
def expand(self, pcoll):
# Transform logic goes here.
return pcoll | beam.Map(lambda x: len(x))
@@ -1420,6 +1447,7 @@ class PTransformTest(unittest.TestCase):
class SlowlyChangingSideInputsTest(unittest.TestCase):
"""Tests for PTransform."""
+
def test_side_input_slow_update(self):
temp_file = tempfile.NamedTemporaryFile(delete=True)
src_file_pattern = temp_file.name
@@ -1441,12 +1469,13 @@ class SlowlyChangingSideInputsTest(unittest.TestCase):
for j in range(count):
f.write('f' + idstr + 'a' + str(j) + '\n')
- sample_main_input_elements = ([first_ts - 2, # no output due to no SI
- first_ts + 1, # First window
- first_ts + 8, # Second window
- first_ts + 15, # Third window
- first_ts + 22, # Fourth window
- ])
+ sample_main_input_elements = ([
+ first_ts - 2, # no output due to no SI
+ first_ts + 1, # First window
+ first_ts + 8, # Second window
+ first_ts + 15, # Third window
+ first_ts + 22, # Fourth window
+ ])
pipeline, pipeline_result = snippets.side_input_slow_update(
src_file_pattern, first_ts, last_ts, interval,
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximatequantiles_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximatequantiles_test.py
index 2adfcd05b99..e8ec7e794cb 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximatequantiles_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximatequantiles_test.py
@@ -32,10 +32,11 @@ from . import approximatequantiles
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch(
'apache_beam.examples.snippets.transforms.aggregation.'
- 'approximatequantiles.print',
- lambda x: x)
+ 'approximatequantiles.print', lambda x: x)
class ApproximateQuantilesTest(unittest.TestCase):
+
def test_approximatequantiles(self):
+
def check_result(quantiles):
assert_that(quantiles, equal_to([[0, 250, 500, 750, 1000]]))
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximateunique_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximateunique_test.py
index c945cec534b..96dd1a58746 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximateunique_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/approximateunique_test.py
@@ -33,10 +33,11 @@ from . import approximateunique
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch(
'apache_beam.examples.snippets.transforms.aggregation.'
- 'approximateunique.print',
- lambda x: x)
+ 'approximateunique.print', lambda x: x)
class ApproximateUniqueTest(unittest.TestCase):
+
def test_approximateunique(self):
+
def check_result(approx_count):
actual_count = 1000
sample_size = 16
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/batchelements_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/batchelements_test.py
index b370f8cd16b..bd29fa291ef 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/batchelements_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/batchelements_test.py
@@ -57,6 +57,7 @@ def identity(x):
identity)
# pylint: enable=line-too-long
class BatchElementsTest(unittest.TestCase):
+
def test_batchelements(self):
batchelements.batchelements(check_batches)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py
index ad4ed99a6e2..e4d31b36a0a 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/cogroupbykey_test.py
@@ -51,6 +51,7 @@ def check_plants(actual):
'apache_beam.examples.snippets.transforms.aggregation.cogroupbykey.print',
str)
class CoGroupByKeyTest(unittest.TestCase):
+
def test_cogroupbykey(self):
cogroupbykey.cogroupbykey(check_plants)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_combinefn.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_combinefn.py
index 18e15e9e9bb..75fae060fa7 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_combinefn.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_combinefn.py
@@ -39,6 +39,7 @@ def combineglobally_combinefn(test=None):
import apache_beam as beam
class PercentagesFn(beam.CombineFn):
+
def create_accumulator(self):
return {}
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py
index e7435c644e2..b8886ba4e1e 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineglobally_test.py
@@ -81,6 +81,7 @@ def check_percentages(actual):
str)
# pylint: enable=line-too-long
class CombineGloballyTest(unittest.TestCase):
+
def test_combineglobally_function(self):
combineglobally_function.combineglobally_function(check_common_items)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_combinefn.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_combinefn.py
index 5b5606401b5..4b9f2d03bb1 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_combinefn.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_combinefn.py
@@ -39,6 +39,7 @@ def combineperkey_combinefn(test=None):
import apache_beam as beam
class AverageFn(beam.CombineFn):
+
def create_accumulator(self):
sum = 0.0
count = 0
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_singleton.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_singleton.py
index b20571bde62..eb618266717 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_singleton.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_side_inputs_singleton.py
@@ -37,8 +37,7 @@ def combineperkey_side_inputs_singleton(test=None):
('🍅', 3),
])
| 'Saturated sum' >> beam.CombinePerKey(
- lambda values,
- max_value: min(sum(values), max_value),
+ lambda values, max_value: min(sum(values), max_value),
max_value=beam.pvalue.AsSingleton(max_value))
| beam.Map(print))
# [END combineperkey_side_inputs_singleton]
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py
index 6492734246a..1162ade2c27 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combineperkey_test.py
@@ -101,6 +101,7 @@ def check_average(actual):
'apache_beam.examples.snippets.transforms.aggregation.combineperkey_combinefn.print',
str)
class CombinePerKeyTest(unittest.TestCase):
+
def test_combineperkey_simple(self):
combineperkey_simple.combineperkey_simple(check_total)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_combinefn.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_combinefn.py
index d32b3342448..0b3028507ba 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_combinefn.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_combinefn.py
@@ -39,6 +39,7 @@ def combinevalues_combinefn(test=None):
import apache_beam as beam
class AverageFn(beam.CombineFn):
+
def create_accumulator(self):
return {}
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_singleton.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_singleton.py
index efec1635c18..c07697b2c88 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_singleton.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_side_inputs_singleton.py
@@ -49,8 +49,7 @@ def combinevalues_side_inputs_singleton(test=None):
('🍅', [4, 5, 3]),
])
| 'Saturated sum' >> beam.CombineValues(
- lambda values,
- max_value: min(sum(values), max_value),
+ lambda values, max_value: min(sum(values), max_value),
max_value=beam.pvalue.AsSingleton(max_value))
| beam.Map(print))
# [END combinevalues_side_inputs_singleton]
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py
index 8693e3f300f..3485306cff6 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/combinevalues_test.py
@@ -95,6 +95,7 @@ def check_percentages_per_season(actual):
'apache_beam.examples.snippets.transforms.aggregation.combinevalues_combinefn.print',
str)
class CombineValuesTest(unittest.TestCase):
+
def test_combinevalues_function(self):
combinevalues_function.combinevalues_function(check_saturated_total)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py
index 59f95a1eec1..cf2b0a03e74 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/count_test.py
@@ -70,6 +70,7 @@ def check_total_unique_elements(actual):
'apache_beam.examples.snippets.transforms.aggregation.count_per_element.print',
str)
class CountTest(unittest.TestCase):
+
def test_count_globally(self):
count_globally.count_globally(check_total_elements)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py
index 6f1246ee363..098a62499d8 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/distinct_test.py
@@ -41,6 +41,7 @@ def check_unique_elements(actual):
@mock.patch(
'apache_beam.examples.snippets.transforms.aggregation.distinct.print', str)
class DistinctTest(unittest.TestCase):
+
def test_distinct(self):
distinct.distinct(check_unique_elements)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py
index 1a9b8f158c1..33f77af1731 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupbykey_test.py
@@ -45,6 +45,7 @@ def check_produce_counts(actual):
'apache_beam.examples.snippets.transforms.aggregation.groupbykey.print',
str)
class GroupByKeyTest(unittest.TestCase):
+
def test_groupbykey(self):
groupbykey.groupbykey(check_produce_counts)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupintobatches_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupintobatches_test.py
index 09449f77f6e..4619c1d68b4 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupintobatches_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/groupintobatches_test.py
@@ -45,6 +45,7 @@ def check_batches_with_keys(actual):
str)
# pylint: enable=line-too-long
class GroupIntoBatchesTest(unittest.TestCase):
+
def test_groupintobatches(self):
groupintobatches.groupintobatches(check_batches_with_keys)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py
index 0661fa86287..a8e4531d21a 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/latest_test.py
@@ -55,6 +55,7 @@ def check_latest_elements_per_key(actual):
'apache_beam.examples.snippets.transforms.aggregation.latest_per_key.print',
str)
class LatestTest(unittest.TestCase):
+
def test_latest_globally(self):
latest_globally.latest_globally(check_latest_element)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py
index 1851d555d28..7d0f04cd649 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/max_test.py
@@ -53,6 +53,7 @@ def check_elements_with_max_value_per_key(actual):
'apache_beam.examples.snippets.transforms.aggregation.max_per_key.print',
str)
class MaxTest(unittest.TestCase):
+
def test_max_globally(self):
beam_max_globally.max_globally(check_max_element)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py
index 7d1420ef0ce..b9336c66a15 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/mean_test.py
@@ -53,6 +53,7 @@ def check_elements_with_mean_value_per_key(actual):
'apache_beam.examples.snippets.transforms.aggregation.mean_per_key.print',
str)
class MeanTest(unittest.TestCase):
+
def test_mean_globally(self):
mean_globally.mean_globally(check_mean_element)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min_test.py
index 65b970800de..182682f0ce1 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/min_test.py
@@ -51,6 +51,7 @@ def check_elements_with_min_value_per_key(actual):
'apache_beam.examples.snippets.transforms.aggregation.min_per_key.print',
str)
class MinTest(unittest.TestCase):
+
def test_min_globally(self):
beam_min_globally.min_globally(check_min_element)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py
index a4ecde5bdab..74716499dd2 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sample_test.py
@@ -57,6 +57,7 @@ def check_samples_per_key(actual):
'apache_beam.examples.snippets.transforms.aggregation.sample_fixed_size_per_key.print',
str)
class SampleTest(unittest.TestCase):
+
def test_sample_fixed_size_globally(self):
sample_fixed_size_globally.sample_fixed_size_globally(check_sample)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py
index 122b8c892d7..d926644e38f 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/sum_test.py
@@ -51,6 +51,7 @@ def check_totals_per_key(actual):
'apache_beam.examples.snippets.transforms.aggregation.sum_per_key.print',
str)
class SumTest(unittest.TestCase):
+
def test_sum_globally(self):
beam_sum_globally.sum_globally(check_total)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/tolist_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/tolist_test.py
index 8f5235000c4..d46c105d6bd 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/tolist_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/tolist_test.py
@@ -39,7 +39,9 @@ def identity(x):
identity)
# pylint: enable=line-too-long
class BatchElementsTest(unittest.TestCase):
+
def test_tolist(self):
+
def check(result):
assert_that(
result
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_of.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_of.py
index 51094a993de..8596e9fcfec 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_of.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_of.py
@@ -48,12 +48,11 @@ def top_of(test=None):
'🌽 Corn',
])
| 'Shortest names' >> beam.combiners.Top.Of(
- 2, # number of elements
- key=len, # optional, defaults to the element itself
+ 2, # number of elements
+ key=len, # optional, defaults to the element itself
reverse=True, # optional, defaults to False (largest/descending)
)
- | beam.Map(print)
- )
+ | beam.Map(print))
# [END top_of]
if test:
test(shortest_elements)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_per_key.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_per_key.py
index 676f10ffc31..246de87f55e 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_per_key.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_per_key.py
@@ -51,12 +51,11 @@ def top_per_key(test=None):
('winter', '🍆 Eggplant'),
])
| 'Shortest names per key' >> beam.combiners.Top.PerKey(
- 2, # number of elements
- key=len, # optional, defaults to the value itself
+ 2, # number of elements
+ key=len, # optional, defaults to the value itself
reverse=True, # optional, defaults to False (largest/descending)
)
- | beam.Map(print)
- )
+ | beam.Map(print))
# [END top_per_key]
if test:
test(shortest_elements_per_key)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py
index e928088f2ed..20b360cfe09 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/aggregation/top_test.py
@@ -101,6 +101,7 @@ def check_shortest_elements_per_key(actual):
'apache_beam.examples.snippets.transforms.aggregation.top_per_key.print',
str)
class TopTest(unittest.TestCase):
+
def test_top_largest(self):
top_largest.top_largest(check_largest_elements)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
index 8a7cdfbe926..50677300089 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py
@@ -62,6 +62,7 @@ Row(entity_id='movie_04', title='The Dark Knight', genres='Action')
@mock.patch('sys.stdout', new_callable=StringIO)
class EnrichmentTest(unittest.TestCase):
+
def test_enrichment_with_bigtable(self, mock_stdout):
enrichment_with_bigtable()
output = mock_stdout.getvalue().splitlines()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_dict.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_dict.py
index 64a4b0aa97c..a969765e4d1 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_dict.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_dict.py
@@ -65,8 +65,7 @@ def filter_side_inputs_dict(test=None):
},
])
| 'Filter plants by duration' >> beam.Filter(
- lambda plant,
- keep_duration: keep_duration[plant['duration']],
+ lambda plant, keep_duration: keep_duration[plant['duration']],
keep_duration=beam.pvalue.AsDict(keep_duration),
)
| beam.Map(print))
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_iter.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_iter.py
index 42043a38c35..9db1c61443e 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_iter.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_iter.py
@@ -65,8 +65,7 @@ def filter_side_inputs_iter(test=None):
},
])
| 'Filter valid plants' >> beam.Filter(
- lambda plant,
- valid_durations: plant['duration'] in valid_durations,
+ lambda plant, valid_durations: plant['duration'] in valid_durations,
valid_durations=beam.pvalue.AsIter(valid_durations),
)
| beam.Map(print))
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_singleton.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_singleton.py
index 5971082becd..34662b4e257 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_singleton.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_side_inputs_singleton.py
@@ -61,8 +61,7 @@ def filter_side_inputs_singleton(test=None):
},
])
| 'Filter perennials' >> beam.Filter(
- lambda plant,
- duration: plant['duration'] == duration,
+ lambda plant, duration: plant['duration'] == duration,
duration=beam.pvalue.AsSingleton(perennial),
)
| beam.Map(print))
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py
index fee9e2f03ea..67aeddad398 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/filter_test.py
@@ -73,6 +73,7 @@ def check_valid_plants(actual):
'apache_beam.examples.snippets.transforms.elementwise.filter_side_inputs_dict.print',
str)
class FilterTest(unittest.TestCase):
+
def test_filter_function(self):
filter_function.filter_function(check_perennials)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_singleton.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_singleton.py
index 84aed00b104..6abf3f485a1 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_singleton.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_side_inputs_singleton.py
@@ -48,8 +48,7 @@ def flatmap_side_inputs_singleton(test=None):
'🍅Tomato,🥔Potato',
])
| 'Split words' >> beam.FlatMap(
- lambda text,
- delimiter: text.split(delimiter),
+ lambda text, delimiter: text.split(delimiter),
delimiter=beam.pvalue.AsSingleton(delimiter),
)
| beam.Map(print))
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py
index 6dd02a208a7..7bd6e74dda1 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/flatmap_test.py
@@ -91,6 +91,7 @@ def check_valid_plants(actual):
'apache_beam.examples.snippets.transforms.elementwise.flatmap_side_inputs_dict.print',
str)
class FlatMapTest(unittest.TestCase):
+
def test_flatmap_simple(self):
flatmap_simple.flatmap_simple(check_plants)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py
index fcfe370234f..6dee2d5f98f 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/keys_test.py
@@ -43,6 +43,7 @@ def check_icons(actual):
@mock.patch(
'apache_beam.examples.snippets.transforms.elementwise.keys.print', str)
class KeysTest(unittest.TestCase):
+
def test_keys(self):
keys.keys(check_icons)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py
index fa494314bd3..549245b4d7d 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/kvswap_test.py
@@ -43,6 +43,7 @@ def check_plants(actual):
@mock.patch(
'apache_beam.examples.snippets.transforms.elementwise.kvswap.print', str)
class KvSwapTest(unittest.TestCase):
+
def test_kvswap(self):
kvswap.kvswap(check_plants)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_context.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_context.py
index 26b5558928c..a827e72b3f3 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_context.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_context.py
@@ -71,9 +71,8 @@ def map_context(test=None):
],
reshuffle=False)
| 'Strip header' >> beam.Map(
- lambda text,
- a=beam.DoFn.SetupContextParam(random_nonce),
- b=beam.DoFn.BundleContextParam(random_nonce): f"{text} {a} {b}")
+ lambda text, a=beam.DoFn.SetupContextParam(random_nonce), b=beam.
+ DoFn.BundleContextParam(random_nonce): f"{text} {a} {b}")
| beam.Map(print))
# [END map_context]
if test:
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_iter.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_iter.py
index c155764a0cd..f0be086e918 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_iter.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_iter.py
@@ -51,8 +51,7 @@ def map_side_inputs_iter(test=None):
'# 🥔Potato\n',
])
| 'Strip header' >> beam.Map(
- lambda text,
- chars: text.strip(''.join(chars)),
+ lambda text, chars: text.strip(''.join(chars)),
chars=beam.pvalue.AsIter(chars),
)
| beam.Map(print))
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_singleton.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_singleton.py
index 323134e315e..8dc6ebb354a 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_singleton.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_side_inputs_singleton.py
@@ -51,8 +51,7 @@ def map_side_inputs_singleton(test=None):
'# 🥔Potato\n',
])
| 'Strip header' >> beam.Map(
- lambda text,
- chars: text.strip(chars),
+ lambda text, chars: text.strip(chars),
chars=beam.pvalue.AsSingleton(chars),
)
| beam.Map(print))
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py
index db766bb8e7c..2b0e2c06aa3 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/map_test.py
@@ -88,6 +88,7 @@ def check_plant_details(actual):
'apache_beam.examples.snippets.transforms.elementwise.map_context.print',
str)
class MapTest(unittest.TestCase):
+
def test_map_simple(self):
map_simple.map_simple(check_plants)
@@ -116,6 +117,7 @@ class MapTest(unittest.TestCase):
import re
def check_nonces(output):
+
def shares_same_nonces(elements):
s = set(re.search(r'\d+ \d+', e).group(0) for e in elements)
assert len(s) == 1, s
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py
index 261b480b108..f046e1bd424 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/mltransform_test.py
@@ -68,6 +68,7 @@ Row(x=array([3]))
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch('sys.stdout', new_callable=StringIO)
class MLTransformStdOutTest(unittest.TestCase):
+
def test_mltransform_compute_and_apply_vocab(self, mock_stdout):
mltransform_compute_and_apply_vocabulary()
predicted = mock_stdout.getvalue().splitlines()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn.py
index 0a6c27a21d8..eebb847c5ee 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn.py
@@ -38,6 +38,7 @@ def pardo_dofn(test=None):
import apache_beam as beam
class SplitWords(beam.DoFn):
+
def __init__(self, delimiter=','):
self.delimiter = delimiter
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_methods.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_methods.py
index 86851960256..cd0f1d8015a 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_methods.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_methods.py
@@ -38,6 +38,7 @@ def pardo_dofn_methods(test=None):
import apache_beam as beam
class DoFnMethods(beam.DoFn):
+
def __init__(self):
print('__init__')
self.window = beam.transforms.window.GlobalWindow()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_params.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_params.py
index ae777836d21..1f78aa13b06 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_params.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_dofn_params.py
@@ -39,6 +39,7 @@ def pardo_dofn_params(test=None):
import apache_beam as beam
class AnalyzeElement(beam.DoFn):
+
def process(
self,
elem,
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py
index 1b72640f5c5..29583438936 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/pardo_test.py
@@ -89,6 +89,7 @@ teardown
'apache_beam.examples.snippets.transforms.elementwise.pardo_dofn_params.print',
str)
class ParDoTest(unittest.TestCase):
+
def test_pardo_dofn(self):
pardo_dofn.pardo_dofn(check_plants)
@@ -99,6 +100,7 @@ class ParDoTest(unittest.TestCase):
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch('sys.stdout', new_callable=StringIO)
class ParDoStdoutTest(unittest.TestCase):
+
def test_pardo_dofn_methods(self, mock_stdout):
expected = pardo_dofn_methods.pardo_dofn_methods(check_dofn_methods)
actual = mock_stdout.getvalue().splitlines()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py
index 3ffc49abe75..80bc614bc1b 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/partition_test.py
@@ -91,6 +91,7 @@ train: {'icon': '🥔', 'name': 'Potato', 'duration': 'perennial'}
'apache_beam.examples.snippets.transforms.elementwise.partition_multiple_arguments.print',
lambda elem: elem)
class PartitionTest(unittest.TestCase):
+
def test_partition_function(self):
partition_function.partition_function(check_partitions)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py
index 2ae53089173..310d7c83f5d 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/regex_test.py
@@ -156,6 +156,7 @@ def check_split(actual):
'apache_beam.examples.snippets.transforms.elementwise.regex_split.print',
str)
class RegexTest(unittest.TestCase):
+
def test_matches(self):
regex_matches.regex_matches(check_matches)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference.py
index 8021a60929a..2a041ae9c2c 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference.py
@@ -23,6 +23,7 @@ import torch
class LinearRegression(torch.nn.Module):
+
def __init__(self, input_dim=1, output_dim=1):
super().__init__()
self.linear = torch.nn.Linear(input_dim, output_dim)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference_test.py
index 8dd46c659f3..90eed2f7b51 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference_test.py
@@ -93,6 +93,7 @@ PredictionResult(example=array([90.], dtype=float32), inference=array([450.], dt
'apache_beam.examples.snippets.transforms.elementwise.runinference_sklearn_keyed_model_handler.print',
str)
class RunInferenceTest(unittest.TestCase):
+
def test_sklearn_unkeyed_model_handler(self):
runinference_sklearn_unkeyed_model_handler.sklearn_unkeyed_model_handler(
check_sklearn_unkeyed_model_handler)
@@ -105,6 +106,7 @@ class RunInferenceTest(unittest.TestCase):
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch('sys.stdout', new_callable=StringIO)
class RunInferenceStdoutTest(unittest.TestCase):
+
@pytest.mark.uses_pytorch
def test_check_torch_keyed_model_handler(self, mock_stdout):
runinference.torch_keyed_model_handler()
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py
index e63282a7573..c263b3e18d0 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/tostring_test.py
@@ -75,6 +75,7 @@ def check_plants_csv(actual):
'apache_beam.examples.snippets.transforms.elementwise.tostring_iterables.print',
str)
class ToStringTest(unittest.TestCase):
+
def test_tostring_kvs(self):
tostring_kvs.tostring_kvs(check_plants)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py
index acd84215a69..77d6d4b3979 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/values_test.py
@@ -43,6 +43,7 @@ Potato
@mock.patch(
'apache_beam.examples.snippets.transforms.elementwise.values.print', str)
class ValuesTest(unittest.TestCase):
+
def test_values(self):
values.values(check_plants)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_event_time.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_event_time.py
index c5e4013e030..3f2cbb5ea32 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_event_time.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_event_time.py
@@ -38,6 +38,7 @@ def withtimestamps_event_time(test=None):
import apache_beam as beam
class GetTimestamp(beam.DoFn):
+
def process(self, plant, timestamp=beam.DoFn.TimestampParam):
yield '{} - {}'.format(timestamp.to_utc_datetime(), plant['name'])
@@ -45,17 +46,26 @@ def withtimestamps_event_time(test=None):
plant_timestamps = (
pipeline
| 'Garden plants' >> beam.Create([
- {'name': 'Strawberry', 'season': 1585699200}, # April, 2020
- {'name': 'Carrot', 'season': 1590969600}, # June, 2020
- {'name': 'Artichoke', 'season': 1583020800}, # March, 2020
- {'name': 'Tomato', 'season': 1588291200}, # May, 2020
- {'name': 'Potato', 'season': 1598918400}, # September, 2020
+ {
+ 'name': 'Strawberry', 'season': 1585699200
+ }, # April, 2020
+ {
+ 'name': 'Carrot', 'season': 1590969600
+ }, # June, 2020
+ {
+ 'name': 'Artichoke', 'season': 1583020800
+ }, # March, 2020
+ {
+ 'name': 'Tomato', 'season': 1588291200
+ }, # May, 2020
+ {
+ 'name': 'Potato', 'season': 1598918400
+ }, # September, 2020
])
| 'With timestamps' >> beam.Map(
lambda plant: beam.window.TimestampedValue(plant, plant['season']))
| 'Get timestamp' >> beam.ParDo(GetTimestamp())
- | beam.Map(print)
- )
+ | beam.Map(print))
# [END withtimestamps_event_time]
if test:
test(plant_timestamps)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_logical_clock.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_logical_clock.py
index 9a2ba36bacd..5ce187d807d 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_logical_clock.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_logical_clock.py
@@ -38,6 +38,7 @@ def withtimestamps_logical_clock(test=None):
import apache_beam as beam
class GetTimestamp(beam.DoFn):
+
def process(self, plant, timestamp=beam.DoFn.TimestampParam):
event_id = int(timestamp.micros / 1e6) # equivalent to seconds
yield '{} - {}'.format(event_id, plant['name'])
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_processing_time.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_processing_time.py
index eacb02d9985..871c16b2bc0 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_processing_time.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_processing_time.py
@@ -39,6 +39,7 @@ def withtimestamps_processing_time(test=None):
import time
class GetTimestamp(beam.DoFn):
+
def process(self, plant, timestamp=beam.DoFn.TimestampParam):
yield '{} - {}'.format(timestamp.to_utc_datetime(), plant['name'])
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py
index 9a86a6f54c5..9184b917e53 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/elementwise/withtimestamps_test.py
@@ -82,6 +82,7 @@ def check_plant_processing_times(actual):
'apache_beam.examples.snippets.transforms.elementwise.withtimestamps_processing_time.print',
str)
class WithTimestampsTest(unittest.TestCase):
+
def test_event_time(self):
withtimestamps_event_time.withtimestamps_event_time(check_plant_timestamps)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/other/create_test.py b/sdks/python/apache_beam/examples/snippets/transforms/other/create_test.py
index 0457fdae777..4ed85b4c9ee 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/other/create_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/other/create_test.py
@@ -43,6 +43,7 @@ def check_create(actual):
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch('apache_beam.examples.snippets.transforms.other.create.print', str)
class CreateTest(unittest.TestCase):
+
def test_create(self):
create.create(check_create)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/other/flatten_test.py b/sdks/python/apache_beam/examples/snippets/transforms/other/flatten_test.py
index bd029c6eeb2..5fe473d5a82 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/other/flatten_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/other/flatten_test.py
@@ -45,6 +45,7 @@ def check_flatten(actual):
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch('apache_beam.examples.snippets.transforms.other.flatten.print', str)
class FlattenTest(unittest.TestCase):
+
def test_flatten(self):
flatten.flatten(check_flatten)
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/other/window.py b/sdks/python/apache_beam/examples/snippets/transforms/other/window.py
index 484917e7765..8f7a1047792 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/other/window.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/other/window.py
@@ -41,22 +41,38 @@ def window(test=None):
import apache_beam as beam
with beam.Pipeline() as pipeline:
- produce = (pipeline
- | 'Garden plants' >> beam.Create([
- {'name': 'Strawberry', 'season': 1585699200}, # April, 2020
- {'name': 'Strawberry', 'season': 1588291200}, # May, 2020
- {'name': 'Carrot', 'season': 1590969600}, # June, 2020
- {'name': 'Artichoke', 'season': 1583020800}, # March, 2020
- {'name': 'Artichoke', 'season': 1585699200}, # April, 2020
- {'name': 'Tomato', 'season': 1588291200}, # May, 2020
- {'name': 'Potato', 'season': 1598918400}, # September, 2020
- ])
- | 'With timestamps' >> beam.Map(lambda plant: beam.window.TimestampedValue(plant['name'], plant['season']))
- | 'Window into fixed 2-month windows' >> beam.WindowInto(
- beam.window.FixedWindows(2 * 30 * 24 * 60 * 60))
- | 'Count per window' >> beam.combiners.Count.PerElement()
- | 'Print results' >> beam.Map(print)
- )
+ produce = (
+ pipeline
+ | 'Garden plants' >> beam.Create([
+ {
+ 'name': 'Strawberry', 'season': 1585699200
+ }, # April, 2020
+ {
+ 'name': 'Strawberry', 'season': 1588291200
+ }, # May, 2020
+ {
+ 'name': 'Carrot', 'season': 1590969600
+ }, # June, 2020
+ {
+ 'name': 'Artichoke', 'season': 1583020800
+ }, # March, 2020
+ {
+ 'name': 'Artichoke', 'season': 1585699200
+ }, # April, 2020
+ {
+ 'name': 'Tomato', 'season': 1588291200
+ }, # May, 2020
+ {
+ 'name': 'Potato', 'season': 1598918400
+ }, # September, 2020
+ ])
+ | 'With timestamps' >> beam.Map(
+ lambda plant: beam.window.TimestampedValue(
+ plant['name'], plant['season']))
+ | 'Window into fixed 2-month windows' >> beam.WindowInto(
+ beam.window.FixedWindows(2 * 30 * 24 * 60 * 60))
+ | 'Count per window' >> beam.combiners.Count.PerElement()
+ | 'Print results' >> beam.Map(print))
# [END window]
if test:
diff --git a/sdks/python/apache_beam/examples/snippets/transforms/other/window_test.py b/sdks/python/apache_beam/examples/snippets/transforms/other/window_test.py
index 6fb53b04b28..a95a4ac80fd 100644
--- a/sdks/python/apache_beam/examples/snippets/transforms/other/window_test.py
+++ b/sdks/python/apache_beam/examples/snippets/transforms/other/window_test.py
@@ -43,6 +43,7 @@ def check_window(actual):
@mock.patch('apache_beam.Pipeline', TestPipeline)
@mock.patch('apache_beam.examples.snippets.transforms.other.window.print', str)
class WindowTest(unittest.TestCase):
+
def test_window(self):
window.window(check_window)
diff --git a/sdks/python/apache_beam/examples/snippets/util.py b/sdks/python/apache_beam/examples/snippets/util.py
index 7911b32ce4d..181e90639ca 100644
--- a/sdks/python/apache_beam/examples/snippets/util.py
+++ b/sdks/python/apache_beam/examples/snippets/util.py
@@ -37,6 +37,7 @@ def assert_matches_stdout(
comparing them. Can be used to sort lists before comparing.
label (str): [optional] Label to make transform names unique.
"""
+
def stdout_to_python_object(elem_str):
try:
elem = ast.literal_eval(elem_str)
diff --git a/sdks/python/apache_beam/examples/snippets/util_test.py b/sdks/python/apache_beam/examples/snippets/util_test.py
index fae26a339d5..857ca100a70 100644
--- a/sdks/python/apache_beam/examples/snippets/util_test.py
+++ b/sdks/python/apache_beam/examples/snippets/util_test.py
@@ -28,6 +28,7 @@ from apache_beam.testing.test_pipeline import TestPipeline
class UtilTest(unittest.TestCase):
+
def test_assert_matches_stdout_object(self):
expected = [
"{'a': '🍓', 'b': True}",
diff --git a/sdks/python/apache_beam/examples/sql_taxi.py b/sdks/python/apache_beam/examples/sql_taxi.py
index e8a29806d72..9cee37305f6 100644
--- a/sdks/python/apache_beam/examples/sql_taxi.py
+++ b/sdks/python/apache_beam/examples/sql_taxi.py
@@ -50,8 +50,8 @@ def run(output_topic, pipeline_args):
# Use beam.Row to create a schema-aware PCollection
| "Create beam Row" >> beam.Map(
lambda x: beam.Row(
- ride_status=str(x['ride_status']),
- passenger_count=int(x['passenger_count'])))
+ ride_status=str(x['ride_status']), passenger_count=int(
+ x['passenger_count'])))
# SqlTransform will computes result within an existing window
| "15s fixed windows" >> beam.WindowInto(beam.window.FixedWindows(15))
# Aggregate drop offs and pick ups that occur within each 15s window
@@ -68,13 +68,10 @@ def run(output_topic, pipeline_args):
# the outputs of the query.
# Collect those attributes, as well as window information, into a dict
| "Assemble Dictionary" >> beam.Map(
- lambda row,
- window=beam.DoFn.WindowParam: {
- "ride_status": row.ride_status,
- "num_rides": row.num_rides,
- "total_passengers": row.total_passengers,
- "window_start": window.start.to_rfc3339(),
- "window_end": window.end.to_rfc3339()
+ lambda row, window=beam.DoFn.WindowParam: {
+ "ride_status": row.ride_status, "num_rides": row.num_rides,
+ "total_passengers": row.total_passengers, "window_start": window
+ .start.to_rfc3339(), "window_end": window.end.to_rfc3339()
})
| "Convert to JSON" >> beam.Map(json.dumps)
| "UTF-8 encode" >> beam.Map(lambda s: s.encode("utf-8"))
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py b/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
index af99a4ab537..fcd074653b8 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_debugging.py
@@ -53,6 +53,7 @@ from apache_beam.transforms.core import ParDo
class PrintFn(beam.DoFn):
"""A DoFn that prints label, element, its window, and its timstamp. """
+
def __init__(self, label):
self.label = label
@@ -75,6 +76,7 @@ class AddTimestampFn(beam.DoFn):
For example, 120 and Sometext will result in:
(120, Timestamp(120) and (Sometext, Timestamp(1234567890).
"""
+
def process(self, element):
logging.info('Adding timestamp to: %s', element)
try:
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_debugging_it_test.py b/sdks/python/apache_beam/examples/streaming_wordcount_debugging_it_test.py
index f3460ec24f1..cfac01c2e12 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_debugging_it_test.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_debugging_it_test.py
@@ -57,6 +57,7 @@ WAIT_UNTIL_FINISH_DURATION = 6 * 60 * 1000 # in milliseconds
class StreamingWordcountDebuggingIT(unittest.TestCase):
+
def setUp(self):
self.test_pipeline = TestPipeline(is_integration_test=True)
self.project = self.test_pipeline.get_option('project')
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py b/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py
index 2ab80993080..98bcd7d9d70 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_debugging_test.py
@@ -39,10 +39,12 @@ except ImportError:
class StreamingWordcountDebugging(unittest.TestCase):
+
@unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
@mock.patch('apache_beam.io.ReadFromPubSub')
@mock.patch('apache_beam.io.WriteToPubSub')
def test_streaming_wordcount_debugging(self, *unused_mocks):
+
def FakeReadFromPubSub(topic=None, subscription=None, values=None):
expected_topic = topic
expected_subscription = subscription
@@ -55,6 +57,7 @@ class StreamingWordcountDebugging(unittest.TestCase):
return _inner
class AssertTransform(beam.PTransform):
+
def __init__(self, matcher):
self.matcher = matcher
diff --git a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
index 3fc4bf0fbc2..4c85c39574c 100644
--- a/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
+++ b/sdks/python/apache_beam/examples/streaming_wordcount_it_test.py
@@ -43,6 +43,7 @@ WAIT_UNTIL_FINISH_DURATION = 10 * 60 * 1000 # in milliseconds
class StreamingWordCountIT(unittest.TestCase):
+
def setUp(self):
self.test_pipeline = TestPipeline(is_integration_test=True)
self.project = self.test_pipeline.get_option('project')
diff --git a/sdks/python/apache_beam/examples/windowed_wordcount.py b/sdks/python/apache_beam/examples/windowed_wordcount.py
index 7889f61b9a6..c4a82cd3b17 100644
--- a/sdks/python/apache_beam/examples/windowed_wordcount.py
+++ b/sdks/python/apache_beam/examples/windowed_wordcount.py
@@ -40,6 +40,7 @@ def find_words(element):
class FormatDoFn(beam.DoFn):
+
def process(self, element, window=beam.DoFn.WindowParam):
ts_format = '%Y-%m-%d %H:%M:%S.%f UTC'
window_start = window.start.to_utc_datetime().strftime(ts_format)
diff --git a/sdks/python/apache_beam/examples/wordcount.py b/sdks/python/apache_beam/examples/wordcount.py
index a9138647581..b4b225e980a 100644
--- a/sdks/python/apache_beam/examples/wordcount.py
+++ b/sdks/python/apache_beam/examples/wordcount.py
@@ -50,6 +50,7 @@ from apache_beam.runners.runner import PipelineResult
class WordExtractingDoFn(beam.DoFn):
"""Parse each line of input text into words."""
+
def process(self, element):
"""Returns an iterator over the words of this element.
diff --git a/sdks/python/apache_beam/examples/wordcount_debugging.py b/sdks/python/apache_beam/examples/wordcount_debugging.py
index 581bbd3adc1..4c11168862e 100644
--- a/sdks/python/apache_beam/examples/wordcount_debugging.py
+++ b/sdks/python/apache_beam/examples/wordcount_debugging.py
@@ -78,6 +78,7 @@ from apache_beam.testing.util import equal_to
class FilterTextFn(beam.DoFn):
"""A DoFn that filters for a specific key based on a regular expression."""
+
def __init__(self, pattern):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
@@ -115,7 +116,9 @@ class CountWords(beam.PTransform):
A PTransform that converts a PCollection containing lines of text into a
PCollection of (word, count) tuples.
"""
+
def expand(self, pcoll):
+
def count_ones(word_ones):
(word, ones) = word_ones
return (word, sum(ones))
diff --git a/sdks/python/apache_beam/examples/wordcount_with_metrics.py b/sdks/python/apache_beam/examples/wordcount_with_metrics.py
index f575a8d7fbb..b7103fb6763 100644
--- a/sdks/python/apache_beam/examples/wordcount_with_metrics.py
+++ b/sdks/python/apache_beam/examples/wordcount_with_metrics.py
@@ -52,6 +52,7 @@ from apache_beam.options.pipeline_options import SetupOptions
class WordExtractingDoFn(beam.DoFn):
"""Parse each line of input text into words."""
+
def __init__(self):
# TODO(BEAM-6158): Revert the workaround once we can pickle super() on py3.
# super().__init__()
diff --git a/sdks/python/apache_beam/examples/wordcount_xlang.py b/sdks/python/apache_beam/examples/wordcount_xlang.py
index 80a0c8fe7a4..823df2ecf34 100644
--- a/sdks/python/apache_beam/examples/wordcount_xlang.py
+++ b/sdks/python/apache_beam/examples/wordcount_xlang.py
@@ -39,6 +39,7 @@ EXPANSION_SERVICE_ADDR = 'localhost:%s' % EXPANSION_SERVICE_PORT
class WordExtractingDoFn(beam.DoFn):
"""Parse each line of input text into words."""
+
def process(self, element):
"""Returns an iterator over the words of this element.
diff --git a/sdks/python/apache_beam/internal/dill_pickler.py b/sdks/python/apache_beam/internal/dill_pickler.py
index e1d6b7e74e4..f96753389f2 100644
--- a/sdks/python/apache_beam/internal/dill_pickler.py
+++ b/sdks/python/apache_beam/internal/dill_pickler.py
@@ -174,6 +174,7 @@ if patch_save_code:
class _NoOpContextManager(object):
+
def __enter__(self):
pass
@@ -251,6 +252,7 @@ def _nested_type_wrapper(fun):
For nested class object only it will save the containing class object so
the nested structure is recreated during unpickle.
"""
+
def wrapper(pickler, obj):
# When the nested class is defined in the __main__ module we do not have to
# do anything special because the pickler itself will save the constituent
@@ -309,7 +311,8 @@ if 'save_module' in dir(dill.dill):
else:
dill_log.info('M2: %s' % obj)
# pylint: disable=protected-access
- pickler.save_reduce(dill.dill._import_module, (obj.__name__, ), obj=obj)
+ pickler.save_reduce(
+ dill.dill._import_module, (obj.__name__, ), obj=obj)
# pylint: enable=protected-access
dill_log.info('# M2')
diff --git a/sdks/python/apache_beam/internal/gcp/auth.py b/sdks/python/apache_beam/internal/gcp/auth.py
index 66c08b8344c..489de12ac4b 100644
--- a/sdks/python/apache_beam/internal/gcp/auth.py
+++ b/sdks/python/apache_beam/internal/gcp/auth.py
@@ -92,6 +92,7 @@ if _GOOGLE_AUTH_AVAILABLE:
upgrading the auth library used by Beam without simultaneously upgrading
all the GCP client libraries (a much larger change).
"""
+
def __init__(self, google_auth_credentials):
self._google_auth_credentials = google_auth_credentials
diff --git a/sdks/python/apache_beam/internal/gcp/auth_test.py b/sdks/python/apache_beam/internal/gcp/auth_test.py
index 654d8e815a5..ef69898ea29 100644
--- a/sdks/python/apache_beam/internal/gcp/auth_test.py
+++ b/sdks/python/apache_beam/internal/gcp/auth_test.py
@@ -31,6 +31,7 @@ except ImportError:
class MockLoggingHandler(logging.Handler):
"""Mock logging handler to check for expected logs."""
+
def __init__(self, *args, **kwargs):
self.reset()
logging.Handler.__init__(self, *args, **kwargs)
@@ -50,6 +51,7 @@ class MockLoggingHandler(logging.Handler):
@unittest.skipIf(gauth is None, 'Google Auth dependencies are not installed')
class AuthTest(unittest.TestCase):
+
@mock.patch('google.auth.default')
def test_auth_with_retrys(self, unused_mock_arg):
pipeline_options = PipelineOptions()
diff --git a/sdks/python/apache_beam/internal/gcp/json_value_test.py b/sdks/python/apache_beam/internal/gcp/json_value_test.py
index 21337de1a10..f28fe97cf93 100644
--- a/sdks/python/apache_beam/internal/gcp/json_value_test.py
+++ b/sdks/python/apache_beam/internal/gcp/json_value_test.py
@@ -37,6 +37,7 @@ except ImportError:
@unittest.skipIf(JsonValue is None, 'GCP dependencies are not installed')
class JsonValueTest(unittest.TestCase):
+
def test_string_to(self):
self.assertEqual(JsonValue(string_value='abc'), to_json_value('abc'))
diff --git a/sdks/python/apache_beam/internal/http_client_test.py b/sdks/python/apache_beam/internal/http_client_test.py
index c8790782a8c..4d40fec5951 100644
--- a/sdks/python/apache_beam/internal/http_client_test.py
+++ b/sdks/python/apache_beam/internal/http_client_test.py
@@ -30,6 +30,7 @@ from apache_beam.internal.http_client import proxy_info_from_environment_var
class HttpClientTest(unittest.TestCase):
+
def test_proxy_from_env_http_with_port(self):
with mock.patch.dict(os.environ, http_proxy='http://localhost:9000'):
proxy_info = proxy_info_from_environment_var('http_proxy')
diff --git a/sdks/python/apache_beam/internal/metrics/cells.py b/sdks/python/apache_beam/internal/metrics/cells.py
index 989dc718304..3f44a4df0e2 100644
--- a/sdks/python/apache_beam/internal/metrics/cells.py
+++ b/sdks/python/apache_beam/internal/metrics/cells.py
@@ -47,6 +47,7 @@ class HistogramCell(MetricCell):
This class is thread safe since underlying histogram object is thread safe.
"""
+
def __init__(self, bucket_type):
self._bucket_type = bucket_type
self.data = HistogramData.identity_element(bucket_type)
@@ -73,6 +74,7 @@ class HistogramCell(MetricCell):
class HistogramCellFactory(MetricCellFactory):
+
def __init__(self, bucket_type):
self._bucket_type = bucket_type
@@ -89,6 +91,7 @@ class HistogramCellFactory(MetricCellFactory):
class HistogramResult(object):
+
def __init__(self, data: 'HistogramData') -> None:
self.data = data
@@ -126,6 +129,7 @@ class HistogramData(object):
This object is not thread safe, so it's not supposed to be modified
outside the HistogramCell.
"""
+
def __init__(self, histogram):
self.histogram = histogram
diff --git a/sdks/python/apache_beam/internal/metrics/cells_test.py b/sdks/python/apache_beam/internal/metrics/cells_test.py
index 066dec4a263..4f1b6bf6854 100644
--- a/sdks/python/apache_beam/internal/metrics/cells_test.py
+++ b/sdks/python/apache_beam/internal/metrics/cells_test.py
@@ -28,6 +28,7 @@ from apache_beam.utils.histogram import LinearBucket
class TestHistogramCell(unittest.TestCase):
+
@classmethod
def _modify_histogram(cls, d):
for i in range(cls.NUM_ITERATIONS):
diff --git a/sdks/python/apache_beam/internal/metrics/metric.py b/sdks/python/apache_beam/internal/metrics/metric.py
index 8acf800ff8c..39fe2143571 100644
--- a/sdks/python/apache_beam/internal/metrics/metric.py
+++ b/sdks/python/apache_beam/internal/metrics/metric.py
@@ -60,6 +60,7 @@ _LOGGER = logging.getLogger(__name__)
class Metrics(object):
+
@staticmethod
def counter(
urn: str,
@@ -106,6 +107,7 @@ class Metrics(object):
class DelegatingHistogram(Histogram):
"""Metrics Histogram that Delegates functionality to MetricsEnvironment."""
+
def __init__(
self,
metric_name: MetricName,
@@ -125,6 +127,7 @@ class Metrics(object):
class MetricLogger(object):
"""Simple object to locally aggregate and log metrics."""
+
def __init__(self) -> None:
self._metric: Dict[MetricName, 'MetricCell'] = {}
self._lock = threading.Lock()
@@ -152,8 +155,8 @@ class MetricLogger(object):
if self._lock.acquire(False):
try:
current_millis = int(time.time() * 1000)
- if ((current_millis - self._last_logging_millis) >
- self.minimum_logging_frequency_msec):
+ if ((current_millis - self._last_logging_millis)
+ > self.minimum_logging_frequency_msec):
logging_metric_info = [
'[Locally aggregated metrics since %s]' %
datetime.datetime.fromtimestamp(
@@ -180,6 +183,7 @@ class ServiceCallMetric(object):
TODO(ajamato): Add Request latency metric.
"""
+
def __init__(
self,
request_count_urn: str,
diff --git a/sdks/python/apache_beam/internal/metrics/metric_test.py b/sdks/python/apache_beam/internal/metrics/metric_test.py
index 22b64ee73ae..2e775036e1b 100644
--- a/sdks/python/apache_beam/internal/metrics/metric_test.py
+++ b/sdks/python/apache_beam/internal/metrics/metric_test.py
@@ -34,6 +34,7 @@ from apache_beam.utils.histogram import LinearBucket
class MetricLoggerTest(unittest.TestCase):
+
@patch('apache_beam.internal.metrics.metric._LOGGER')
def test_log_metrics(self, mock_logger):
logger = MetricLogger()
@@ -44,6 +45,7 @@ class MetricLoggerTest(unittest.TestCase):
logger.log_metrics()
class Contains(str):
+
def __eq__(self, other):
return self in other
@@ -52,6 +54,7 @@ class MetricLoggerTest(unittest.TestCase):
class MetricsTest(unittest.TestCase):
+
def test_create_process_wide(self):
sampler = statesampler.StateSampler('', counters.CounterFactory())
statesampler.set_current_tracker(sampler)
diff --git a/sdks/python/apache_beam/internal/module_test.py b/sdks/python/apache_beam/internal/module_test.py
index ff0ad0c564e..52b387725e9 100644
--- a/sdks/python/apache_beam/internal/module_test.py
+++ b/sdks/python/apache_beam/internal/module_test.py
@@ -25,12 +25,16 @@ from typing import Any
class TopClass(object):
+
class NestedClass(object):
+
def __init__(self, datum):
self.datum = 'X:%s' % datum
class MiddleClass(object):
+
class NestedClass(object):
+
def __init__(self, datum):
self.datum = 'Y:%s' % datum
@@ -45,13 +49,16 @@ def get_lambda_with_closure(message):
class Xyz(object):
"""A class to be pickled."""
+
def foo(self, s):
return re.findall(r'\w+', s)
def create_class(datum):
"""Creates an unnamable class to be pickled."""
+
class Z(object):
+
def get(self):
return 'Z:%s' % datum
diff --git a/sdks/python/apache_beam/internal/util.py b/sdks/python/apache_beam/internal/util.py
index 85a6e4c43b8..bac34bd8286 100644
--- a/sdks/python/apache_beam/internal/util.py
+++ b/sdks/python/apache_beam/internal/util.py
@@ -51,6 +51,7 @@ class ArgumentPlaceholder(object):
Fn object by the time it executes will have such values replaced with real
computed values.
"""
+
def __eq__(self, other):
"""Tests for equality of two placeholder objects.
@@ -97,8 +98,7 @@ def remove_objects_from_args(
# by sorting the entries first. This will be important when putting back
# PValues.
new_kwargs = dict((k, swapper(v)) if isinstance(v, pvalue_class) else (k, v)
- for k,
- v in sorted(kwargs.items()))
+ for k, v in sorted(kwargs.items()))
return (new_args, new_kwargs, pvals)
@@ -123,8 +123,8 @@ def insert_values_in_args(args, kwargs, values):
for arg in args
]
new_kwargs = dict(
- (k, next(v_iter)) if isinstance(v, ArgumentPlaceholder) else (k, v) for k,
- v in sorted(kwargs.items()))
+ (k, next(v_iter)) if isinstance(v, ArgumentPlaceholder) else (k, v)
+ for k, v in sorted(kwargs.items()))
return (new_args, new_kwargs)
diff --git a/sdks/python/apache_beam/internal/util_test.py b/sdks/python/apache_beam/internal/util_test.py
index ded1190f840..2449b5412c3 100644
--- a/sdks/python/apache_beam/internal/util_test.py
+++ b/sdks/python/apache_beam/internal/util_test.py
@@ -26,6 +26,7 @@ from apache_beam.internal.util import remove_objects_from_args
class UtilTest(unittest.TestCase):
+
def test_remove_objects_from_args(self):
args, kwargs, objs = remove_objects_from_args(
[1, 'a'], {'x': 1, 'y': 3.14}, (str, float))
diff --git a/sdks/python/apache_beam/io/avroio.py b/sdks/python/apache_beam/io/avroio.py
index 8b7958a00b8..3b8bd11ffec 100644
--- a/sdks/python/apache_beam/io/avroio.py
+++ b/sdks/python/apache_beam/io/avroio.py
@@ -85,6 +85,7 @@ class ReadFromAvro(PTransform):
that comply with the schema contained in the Avro file that contains those
records.
"""
+
def __init__(
self,
file_pattern=None,
@@ -287,6 +288,7 @@ class ReadAllFromAvroContinuously(ReadAllFromAvro):
class _AvroUtils(object):
+
@staticmethod
def advance_file_past_next_sync_marker(f, sync_marker):
buf_size = 10000
@@ -322,6 +324,7 @@ class _FastAvroSource(filebasedsource.FileBasedSource):
TODO: remove ``_AvroSource`` in favor of using ``_FastAvroSource``
everywhere once it has been more widely tested
"""
+
def read_records(self, file_name, range_tracker):
next_block_start = -1
@@ -367,6 +370,7 @@ class WriteToAvro(beam.transforms.PTransform):
If the input has a schema, a corresponding avro schema will be automatically
generated and used to write the output records."""
+
def __init__(
self,
file_path_prefix,
@@ -411,13 +415,8 @@ class WriteToAvro(beam.transforms.PTransform):
"""
self._schema = schema
self._sink_provider = lambda avro_schema: _create_avro_sink(
- file_path_prefix,
- avro_schema,
- codec,
- file_name_suffix,
- num_shards,
- shard_name_template,
- mime_type)
+ file_path_prefix, avro_schema, codec, file_name_suffix, num_shards,
+ shard_name_template, mime_type)
def expand(self, pcoll):
if self._schema:
@@ -465,6 +464,7 @@ def _create_avro_sink(
class _BaseAvroSink(filebasedsink.FileBasedSink):
"""A base for a sink for avro files. """
+
def __init__(
self,
file_path_prefix,
@@ -496,6 +496,7 @@ class _BaseAvroSink(filebasedsink.FileBasedSink):
class _FastAvroSink(_BaseAvroSink):
"""A sink for avro files using FastAvro. """
+
def __init__(
self,
file_path_prefix,
diff --git a/sdks/python/apache_beam/io/avroio_test.py b/sdks/python/apache_beam/io/avroio_test.py
index 2d25010da48..ad6ccfeb26a 100644
--- a/sdks/python/apache_beam/io/avroio_test.py
+++ b/sdks/python/apache_beam/io/avroio_test.py
@@ -174,8 +174,8 @@ class AvroBase(object):
@pytest.mark.xlang_sql_expansion_service
@unittest.skipIf(
- TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is
- None,
+ TestPipeline().get_pipeline_options().view_as(StandardOptions).runner
+ is None,
"Must be run with a runner that supports staging java artifacts.")
def test_avro_schema_to_beam_schema_with_nullable_atomic_fields(self):
records = []
@@ -396,6 +396,7 @@ class AvroBase(object):
self._run_avro_test(pattern, 100, True, expected_result)
def test_dynamic_work_rebalancing_exhaustive(self):
+
def compare_split_points(file_name):
source = _FastAvroSource(file_name)
splits = [
@@ -625,6 +626,7 @@ class AvroBase(object):
class TestFastAvro(AvroBase, unittest.TestCase):
+
def __init__(self, methodName='runTest'):
super().__init__(methodName)
self.SCHEMA = parse_schema(json.loads(self.SCHEMA_STRING))
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py b/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py
index aee24ac2c05..4706f52f31a 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/boto3_client.py
@@ -40,6 +40,7 @@ class Client(object):
"""
Wrapper for boto3 library
"""
+
def __init__(self, options):
assert boto3 is not None, 'Missing boto3 requirement'
if isinstance(options, pipeline_options.PipelineOptions):
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/client_test.py b/sdks/python/apache_beam/io/aws/clients/s3/client_test.py
index 67797b1736b..b809c7e456b 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/client_test.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/client_test.py
@@ -27,6 +27,7 @@ from apache_beam.options import pipeline_options
class ClientErrorTest(unittest.TestCase):
+
def setUp(self):
# These tests can be run locally against a mock S3 client, or as integration
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py b/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py
index e2d34a11a46..a36de57d308 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/fake_client.py
@@ -26,6 +26,7 @@ from apache_beam.io.aws.clients.s3 import messages
class FakeFile(object):
+
def __init__(self, bucket, key, contents, etag=None):
self.bucket = bucket
self.key = key
@@ -53,6 +54,7 @@ class FakeFile(object):
class FakeS3Client(object):
+
def __init__(self):
self.files = {}
self.list_continuation_tokens = {}
diff --git a/sdks/python/apache_beam/io/aws/clients/s3/messages.py b/sdks/python/apache_beam/io/aws/clients/s3/messages.py
index 8555364b727..902ef36ccd7 100644
--- a/sdks/python/apache_beam/io/aws/clients/s3/messages.py
+++ b/sdks/python/apache_beam/io/aws/clients/s3/messages.py
@@ -22,6 +22,7 @@ class GetRequest():
"""
S3 request object for `Get` command
"""
+
def __init__(self, bucket, object):
self.bucket = bucket
self.object = object
@@ -31,6 +32,7 @@ class UploadResponse():
"""
S3 response object for `StartUpload` command
"""
+
def __init__(self, upload_id):
self.upload_id = upload_id
@@ -39,6 +41,7 @@ class UploadRequest():
"""
S3 request object for `StartUpload` command
"""
+
def __init__(self, bucket, object, mime_type):
self.bucket = bucket
self.object = object
@@ -49,6 +52,7 @@ class UploadPartRequest():
"""
S3 request object for `UploadPart` command
"""
+
def __init__(self, bucket, object, upload_id, part_number, bytes):
self.bucket = bucket
self.object = object
@@ -62,6 +66,7 @@ class UploadPartResponse():
"""
S3 response object for `UploadPart` command
"""
+
def __init__(self, etag, part_number):
self.etag = etag
self.part_number = part_number
@@ -71,6 +76,7 @@ class CompleteMultipartUploadRequest():
"""
S3 request object for `UploadPart` command
"""
+
def __init__(self, bucket, object, upload_id, parts):
# parts is a list of objects of the form
# {'ETag': response.etag, 'PartNumber': response.part_number}
@@ -85,6 +91,7 @@ class ListRequest():
"""
S3 request object for `List` command
"""
+
def __init__(self, bucket, prefix, continuation_token=None):
self.bucket = bucket
self.prefix = prefix
@@ -95,6 +102,7 @@ class ListResponse():
"""
S3 response object for `List` command
"""
+
def __init__(self, items, next_token=None):
self.items = items
self.next_token = next_token
@@ -104,6 +112,7 @@ class Item():
"""
An item in S3
"""
+
def __init__(self, etag, key, last_modified, size, mime_type=None):
self.etag = etag
self.key = key
@@ -116,12 +125,14 @@ class DeleteRequest():
"""
S3 request object for `Delete` command
"""
+
def __init__(self, bucket, object):
self.bucket = bucket
self.object = object
class DeleteBatchRequest():
+
def __init__(self, bucket, objects):
# `objects` is a list of strings corresponding to the keys to be deleted
# in the bucket
@@ -130,6 +141,7 @@ class DeleteBatchRequest():
class DeleteBatchResponse():
+
def __init__(self, deleted, failed, errors):
# `deleted` is a list of strings corresponding to the keys that were deleted
# `failed` is a list of strings corresponding to the keys that caused errors
@@ -140,6 +152,7 @@ class DeleteBatchResponse():
class CopyRequest():
+
def __init__(self, src_bucket, src_key, dest_bucket, dest_key):
self.src_bucket = src_bucket
self.src_key = src_key
@@ -148,6 +161,7 @@ class CopyRequest():
class S3ClientError(Exception):
+
def __init__(self, message=None, code=None):
self.message = message
self.code = code
diff --git a/sdks/python/apache_beam/io/aws/s3filesystem_test.py b/sdks/python/apache_beam/io/aws/s3filesystem_test.py
index 036727cd7a7..07abe3486c7 100644
--- a/sdks/python/apache_beam/io/aws/s3filesystem_test.py
+++ b/sdks/python/apache_beam/io/aws/s3filesystem_test.py
@@ -41,6 +41,7 @@ except ImportError:
@unittest.skipIf(s3filesystem is None, 'AWS dependencies are not installed')
class S3FileSystemTest(unittest.TestCase):
+
def setUp(self):
pipeline_options = PipelineOptions()
self.fs = s3filesystem.S3FileSystem(pipeline_options=pipeline_options)
diff --git a/sdks/python/apache_beam/io/aws/s3io.py b/sdks/python/apache_beam/io/aws/s3io.py
index 887bb4c7baa..01be057ea9c 100644
--- a/sdks/python/apache_beam/io/aws/s3io.py
+++ b/sdks/python/apache_beam/io/aws/s3io.py
@@ -56,6 +56,7 @@ def parse_s3_path(s3_path, object_optional=False):
class S3IO(object):
"""S3 I/O client."""
+
def __init__(self, client=None, options=None):
if client is None and options is None:
raise ValueError('Must provide one of client or options')
@@ -566,6 +567,7 @@ class S3IO(object):
class S3Downloader(Downloader):
+
def __init__(self, client, path, buffer_size):
self._client = client
self._path = path
@@ -602,6 +604,7 @@ class S3Downloader(Downloader):
class S3Uploader(Uploader):
+
def __init__(self, client, path, mime_type='application/octet-stream'):
self._client = client
self._path = path
diff --git a/sdks/python/apache_beam/io/aws/s3io_test.py b/sdks/python/apache_beam/io/aws/s3io_test.py
index ffab9572707..339d7378b88 100644
--- a/sdks/python/apache_beam/io/aws/s3io_test.py
+++ b/sdks/python/apache_beam/io/aws/s3io_test.py
@@ -64,6 +64,7 @@ class TestS3PathParser(unittest.TestCase):
class TestS3IO(unittest.TestCase):
+
def _insert_random_file(self, client, path, size):
bucket, name = s3io.parse_s3_path(path)
contents = os.urandom(size)
diff --git a/sdks/python/apache_beam/io/azure/blobstoragefilesystem_test.py b/sdks/python/apache_beam/io/azure/blobstoragefilesystem_test.py
index c3418e137e8..1ff28793b96 100644
--- a/sdks/python/apache_beam/io/azure/blobstoragefilesystem_test.py
+++ b/sdks/python/apache_beam/io/azure/blobstoragefilesystem_test.py
@@ -42,6 +42,7 @@ except ImportError:
@unittest.skipIf(
blobstoragefilesystem is None, 'Azure dependencies are not installed')
class BlobStorageFileSystemTest(unittest.TestCase):
+
def setUp(self):
pipeline_options = PipelineOptions()
self.fs = blobstoragefilesystem.BlobStorageFileSystem(
diff --git a/sdks/python/apache_beam/io/azure/blobstorageio.py b/sdks/python/apache_beam/io/azure/blobstorageio.py
index cfa4fe7d291..e681b668cc4 100644
--- a/sdks/python/apache_beam/io/azure/blobstorageio.py
+++ b/sdks/python/apache_beam/io/azure/blobstorageio.py
@@ -86,6 +86,7 @@ def get_azfs_url(storage_account, container, blob=''):
class Blob():
"""A Blob in Azure Blob Storage."""
+
def __init__(self, etag, name, last_updated, size, mime_type):
self.etag = etag
self.name = name
@@ -101,6 +102,7 @@ class BlobStorageIOError(IOError, retry.PermanentException):
class BlobStorageError(Exception):
"""Blob Storage client error."""
+
def __init__(self, message=None, code=None):
self.message = message
self.code = code
@@ -108,6 +110,7 @@ class BlobStorageError(Exception):
class BlobStorageIO(object):
"""Azure Blob Storage I/O client."""
+
def __init__(self, client=None, pipeline_options=None):
if client is None:
azure_options = pipeline_options.view_as(AzureOptions)
@@ -654,6 +657,7 @@ class BlobStorageIO(object):
class BlobStorageDownloader(Downloader):
+
def __init__(self, client, path, buffer_size):
self._client = client
self._path = path
@@ -692,6 +696,7 @@ class BlobStorageDownloader(Downloader):
class BlobStorageUploader(Uploader):
+
def __init__(self, client, path, mime_type='application/octet-stream'):
self._client = client
self._path = path
diff --git a/sdks/python/apache_beam/io/components/util.py b/sdks/python/apache_beam/io/components/util.py
index 0726ecb1382..c78d462ba0a 100644
--- a/sdks/python/apache_beam/io/components/util.py
+++ b/sdks/python/apache_beam/io/components/util.py
@@ -34,6 +34,7 @@ class MovingSum(object):
convenience we expose the count of entries as well so this doubles as a
moving average tracker.
"""
+
def __init__(self, window_ms, bucket_ms):
if window_ms < bucket_ms or bucket_ms <= 0:
raise ValueError("window_ms >= bucket_ms > 0 please")
diff --git a/sdks/python/apache_beam/io/concat_source.py b/sdks/python/apache_beam/io/concat_source.py
index 35ae6fe2731..bfed92d2e36 100644
--- a/sdks/python/apache_beam/io/concat_source.py
+++ b/sdks/python/apache_beam/io/concat_source.py
@@ -35,6 +35,7 @@ class ConcatSource(iobase.BoundedSource):
Primarily for internal use, use the ``apache_beam.Flatten`` transform
to create the union of several reads.
"""
+
def __init__(self, sources):
self._source_bundles = [
source if isinstance(source, iobase.SourceBundle) else
@@ -98,6 +99,7 @@ class ConcatRangeTracker(iobase.RangeTracker):
"""For internal use only; no backwards-compatibility guarantees.
Range tracker for ConcatSource"""
+
def __init__(self, start, end, source_bundles):
"""Initializes ``ConcatRangeTracker``
diff --git a/sdks/python/apache_beam/io/concat_source_test.py b/sdks/python/apache_beam/io/concat_source_test.py
index efa24b3975f..38aaa6e6a31 100644
--- a/sdks/python/apache_beam/io/concat_source_test.py
+++ b/sdks/python/apache_beam/io/concat_source_test.py
@@ -85,6 +85,7 @@ class RangeSource(iobase.BoundedSource):
class ConcatSourceTest(unittest.TestCase):
+
def test_range_source(self):
source_test_utils.assert_split_at_fraction_exhaustive(RangeSource(0, 10, 3))
diff --git a/sdks/python/apache_beam/io/debezium.py b/sdks/python/apache_beam/io/debezium.py
index ada25760c27..d94e72b42e4 100644
--- a/sdks/python/apache_beam/io/debezium.py
+++ b/sdks/python/apache_beam/io/debezium.py
@@ -115,6 +115,7 @@ ReadFromDebeziumSchema = NamedTuple(
class _JsonStringToDictionaries(DoFn):
""" A DoFn that consumes a JSON string and yields a python dictionary """
+
def process(self, json_string):
obj = json.loads(json_string)
yield obj
diff --git a/sdks/python/apache_beam/io/external/generate_sequence_test.py b/sdks/python/apache_beam/io/external/generate_sequence_test.py
index d44fd363b95..21246a7d81a 100644
--- a/sdks/python/apache_beam/io/external/generate_sequence_test.py
+++ b/sdks/python/apache_beam/io/external/generate_sequence_test.py
@@ -40,6 +40,7 @@ from apache_beam.transforms.external import JavaExternalTransform
os.environ.get('EXPANSION_PORT'),
"EXPANSION_PORT environment var is not provided.")
class XlangGenerateSequenceTest(unittest.TestCase):
+
def test_generate_sequence(self):
port = os.environ.get('EXPANSION_PORT')
address = 'localhost:%s' % port
diff --git a/sdks/python/apache_beam/io/external/xlang_debeziumio_it_test.py b/sdks/python/apache_beam/io/external/xlang_debeziumio_it_test.py
index abe9530787e..1cbbf070949 100644
--- a/sdks/python/apache_beam/io/external/xlang_debeziumio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_debeziumio_it_test.py
@@ -37,10 +37,11 @@ NUM_RECORDS = 1
@unittest.skipIf(
PostgresContainer is None, 'testcontainers package is not installed')
@unittest.skipIf(
- TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is
- None,
+ TestPipeline().get_pipeline_options().view_as(StandardOptions).runner
+ is None,
'Do not run this test on precommit suites.')
class CrossLanguageDebeziumIOTest(unittest.TestCase):
+
def setUp(self):
self.username = 'debezium'
self.password = 'dbz'
diff --git a/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py b/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py
index 38a405c2d33..8d8fa3c36ff 100644
--- a/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py
@@ -71,8 +71,8 @@ coders.registry.register_coder(JdbcTestRow, coders.RowCoder)
@unittest.skipIf(
PostgresContainer is None, 'testcontainers package is not installed')
@unittest.skipIf(
- TestPipeline().get_pipeline_options().view_as(StandardOptions).runner is
- None,
+ TestPipeline().get_pipeline_options().view_as(StandardOptions).runner
+ is None,
'Do not run this test on precommit suites.')
class CrossLanguageJdbcIOTest(unittest.TestCase):
DbData = typing.NamedTuple(
diff --git a/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py b/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
index a7bf686d064..f9a7b5a9fea 100644
--- a/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py
@@ -66,6 +66,7 @@ class CollectingFn(beam.DoFn):
class CrossLanguageKafkaIO(object):
+
def __init__(
self, bootstrap_servers, topic, null_key, expansion_service=None):
self.bootstrap_servers = bootstrap_servers
@@ -113,6 +114,7 @@ class CrossLanguageKafkaIO(object):
class CrossLanguageKafkaIOTest(unittest.TestCase):
+
@unittest.skipUnless(
os.environ.get('LOCAL_KAFKA_JAR'),
"LOCAL_KAFKA_JAR environment var is not provided.")
diff --git a/sdks/python/apache_beam/io/external/xlang_kafkaio_perf_test.py b/sdks/python/apache_beam/io/external/xlang_kafkaio_perf_test.py
index 08a6baee468..2cb13f88681 100644
--- a/sdks/python/apache_beam/io/external/xlang_kafkaio_perf_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_kafkaio_perf_test.py
@@ -37,6 +37,7 @@ _LOGGER = logging.getLogger(__name__)
class KafkaIOTestOptions(LoadTestOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
parser.add_argument(
@@ -57,6 +58,7 @@ class KafkaIOTestOptions(LoadTestOptions):
class KafkaIOPerfTest:
"""Performance test for cross-language Kafka IO pipeline."""
+
def run(self):
write_test = _KafkaIOBatchWritePerfTest()
read_test = _KafkaIOSDFReadPerfTest()
@@ -65,6 +67,7 @@ class KafkaIOPerfTest:
class _KafkaIOBatchWritePerfTest(LoadTest):
+
def __init__(self):
super().__init__(WRITE_NAMESPACE)
self.test_options = self.pipeline.get_pipeline_options().view_as(
@@ -93,6 +96,7 @@ class _KafkaIOBatchWritePerfTest(LoadTest):
class _KafkaIOSDFReadPerfTest(LoadTest):
+
def __init__(self):
super().__init__(READ_NAMESPACE)
self.test_options = self.pipeline.get_pipeline_options().view_as(
diff --git a/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py b/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py
index c9181fb2a72..007fad31bf0 100644
--- a/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_kinesisio_it_test.py
@@ -79,6 +79,7 @@ RECORD = b'record' + str(uuid.uuid4()).encode()
TestPipeline().get_pipeline_options().view_as(StandardOptions).runner,
'Do not run this test on precommit suites.')
class CrossLanguageKinesisIOTest(unittest.TestCase):
+
@unittest.skipUnless(
TestPipeline().get_option('aws_kinesis_stream'),
'Cannot test on real aws without pipeline options provided')
@@ -242,6 +243,7 @@ class CrossLanguageKinesisIOTest(unittest.TestCase):
class KinesisHelper:
+
def __init__(self, access_key, secret_key, region, service_endpoint):
self.kinesis_client = boto3.client(
service_name='kinesis',
diff --git a/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py b/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py
index f78175a8696..e702edd391f 100644
--- a/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py
+++ b/sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py
@@ -92,11 +92,13 @@ NUM_RECORDS = 100
TestPipeline().get_option('server_name') is None,
'Snowflake IT test requires external configuration to be run.')
class SnowflakeTest(unittest.TestCase):
+
def test_snowflake_write_read(self):
self.run_write()
self.run_read()
def run_write(self):
+
def user_data_mapper(test_row):
return [
str(test_row.number_column).encode('utf-8'),
@@ -137,6 +139,7 @@ class SnowflakeTest(unittest.TestCase):
))
def run_read(self):
+
def csv_mapper(bytes_array):
return TestRow(
int(bytes_array[0]),
diff --git a/sdks/python/apache_beam/io/filebasedio_perf_test.py b/sdks/python/apache_beam/io/filebasedio_perf_test.py
index 78a390d9bed..e3c7965e7f7 100644
--- a/sdks/python/apache_beam/io/filebasedio_perf_test.py
+++ b/sdks/python/apache_beam/io/filebasedio_perf_test.py
@@ -45,6 +45,7 @@ _LOGGER = logging.getLogger(__name__)
class FileBasedIOTestOptions(LoadTestOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
parser.add_argument(
@@ -79,6 +80,7 @@ class SyntheticRecordToStrFn(beam.DoFn):
Output length = 4(ceil[len(key)/3] + ceil[len(value)/3]) + 1
"""
+
def process(self, element):
import base64
yield base64.b64encode(element[0]) + b' ' + base64.b64encode(element[1])
@@ -86,6 +88,7 @@ class SyntheticRecordToStrFn(beam.DoFn):
class CreateFolderFn(beam.DoFn):
"""Create folder at pipeline runtime."""
+
def __init__(self, folder):
self.folder = folder
@@ -97,6 +100,7 @@ class CreateFolderFn(beam.DoFn):
class TextIOPerfTest:
+
def run(self):
write_test = _TextIOWritePerfTest(need_cleanup=False)
read_test = _TextIOReadPerfTest(input_folder=write_test.output_folder)
@@ -105,6 +109,7 @@ class TextIOPerfTest:
class _TextIOWritePerfTest(LoadTest):
+
def __init__(self, need_cleanup=True):
super().__init__(WRITE_NAMESPACE)
self.need_cleanup = need_cleanup
@@ -146,6 +151,7 @@ class _TextIOWritePerfTest(LoadTest):
class _TextIOReadPerfTest(LoadTest):
+
def __init__(self, input_folder):
super().__init__(READ_NAMESPACE)
self.test_options = self.pipeline.get_pipeline_options().view_as(
diff --git a/sdks/python/apache_beam/io/filebasedsink.py b/sdks/python/apache_beam/io/filebasedsink.py
index eb433bd6058..faf05e89aa4 100644
--- a/sdks/python/apache_beam/io/filebasedsink.py
+++ b/sdks/python/apache_beam/io/filebasedsink.py
@@ -427,6 +427,7 @@ class FileBasedSink(iobase.Sink):
class FileBasedSinkWriter(iobase.Writer):
"""The writer for FileBasedSink.
"""
+
def __init__(self, sink, temp_shard_path):
self.sink = sink
self.temp_shard_path = temp_shard_path
@@ -440,10 +441,10 @@ class FileBasedSinkWriter(iobase.Writer):
def at_capacity(self):
return (
self.sink.max_records_per_shard and
- self.num_records_written >= self.sink.max_records_per_shard
- ) or (
- self.sink.max_bytes_per_shard and
- self.sink.byte_counter.bytes_written >= self.sink.max_bytes_per_shard)
+ self.num_records_written >= self.sink.max_records_per_shard) or (
+ self.sink.max_bytes_per_shard and
+ self.sink.byte_counter.bytes_written
+ >= self.sink.max_bytes_per_shard)
def close(self):
self.sink.close(self.temp_handle)
@@ -451,6 +452,7 @@ class FileBasedSinkWriter(iobase.Writer):
class _ByteCountingWriter:
+
def __init__(self, writer):
self.writer = writer
self.bytes_written = 0
diff --git a/sdks/python/apache_beam/io/filebasedsink_test.py b/sdks/python/apache_beam/io/filebasedsink_test.py
index 121bc479200..ae0c9c94ccd 100644
--- a/sdks/python/apache_beam/io/filebasedsink_test.py
+++ b/sdks/python/apache_beam/io/filebasedsink_test.py
@@ -50,6 +50,7 @@ class _TestCaseWithTempDirCleanUp(unittest.TestCase):
Inherited test cases will call self._new_tempdir() to start a temporary dir
which will be deleted at the end of the tests (when tearDown() is called).
"""
+
def setUp(self):
self._tempdirs = []
@@ -79,6 +80,7 @@ class _TestCaseWithTempDirCleanUp(unittest.TestCase):
class MyFileBasedSink(filebasedsink.FileBasedSink):
+
def open(self, temp_path):
# TODO: Fix main session pickling.
# file_handle = super().open(temp_path)
@@ -99,6 +101,7 @@ class MyFileBasedSink(filebasedsink.FileBasedSink):
class TestFileBasedSink(_TestCaseWithTempDirCleanUp):
+
def _common_init(self, sink):
# Manually invoke the generic Sink API.
init_token = sink.initialize_write()
@@ -202,6 +205,7 @@ class TestFileBasedSink(_TestCaseWithTempDirCleanUp):
dir_root_path,
prefix,
separator):
+
def _get_temp_dir(file_path_prefix):
sink = MyFileBasedSink(
file_path_prefix,
diff --git a/sdks/python/apache_beam/io/filebasedsource.py b/sdks/python/apache_beam/io/filebasedsource.py
index 49b1b1d125f..d6bbf8e8d8f 100644
--- a/sdks/python/apache_beam/io/filebasedsource.py
+++ b/sdks/python/apache_beam/io/filebasedsource.py
@@ -241,6 +241,7 @@ def _determine_splittability_from_compression_type(file_path, compression_type):
class _SingleFileSource(iobase.BoundedSource):
"""Denotes a source for a specific file type."""
+
def __init__(
self,
file_based_source,
@@ -339,6 +340,7 @@ class _SingleFileSource(iobase.BoundedSource):
class _ExpandIntoRanges(DoFn):
+
def __init__(
self, splittable, compression_type, desired_bundle_size, min_bundle_size):
self._desired_bundle_size = desired_bundle_size
@@ -372,6 +374,7 @@ class _ExpandIntoRanges(DoFn):
class _ReadRange(DoFn):
+
def __init__(
self,
source_from_file: Union[str, iobase.BoundedSource],
@@ -406,6 +409,7 @@ class ReadAllFiles(PTransform):
PTransform authors who wishes to implement file-based Read transforms that
read a PCollection of files.
"""
+
def __init__(
self,
splittable: bool,
diff --git a/sdks/python/apache_beam/io/filebasedsource_test.py b/sdks/python/apache_beam/io/filebasedsource_test.py
index e68d2afbac9..cb3b245112f 100644
--- a/sdks/python/apache_beam/io/filebasedsource_test.py
+++ b/sdks/python/apache_beam/io/filebasedsource_test.py
@@ -47,6 +47,7 @@ from apache_beam.transforms.display_test import DisplayDataItemMatcher
class LineSource(FileBasedSource):
+
def read_records(self, file_name, range_tracker):
f = self.open_file(file_name)
try:
@@ -178,7 +179,9 @@ def write_pattern(lines_per_file, no_data=False):
class TestConcatSource(unittest.TestCase):
+
class DummySource(iobase.BoundedSource):
+
def __init__(self, values):
self._values = values
@@ -252,6 +255,7 @@ class TestConcatSource(unittest.TestCase):
class TestFileBasedSource(unittest.TestCase):
+
def setUp(self):
# Reducing the size of thread pools. Without this test execution may fail in
# environments with limited amount of resources.
@@ -584,10 +588,12 @@ class TestFileBasedSource(unittest.TestCase):
assert_that(pcoll, equal_to(lines))
def test_splits_get_coder_from_fbs(self):
+
class DummyCoder(object):
val = 12345
class FileBasedSourceWithCoder(LineSource):
+
def default_output_coder(self):
return DummyCoder()
@@ -601,6 +607,7 @@ class TestFileBasedSource(unittest.TestCase):
class TestSingleFileSource(unittest.TestCase):
+
def setUp(self):
# Reducing the size of thread pools. Without this test execution may fail in
# environments with limited amount of resources.
diff --git a/sdks/python/apache_beam/io/fileio.py b/sdks/python/apache_beam/io/fileio.py
index 111206a18a2..4d33cb644e4 100644
--- a/sdks/python/apache_beam/io/fileio.py
+++ b/sdks/python/apache_beam/io/fileio.py
@@ -161,6 +161,7 @@ class EmptyMatchTreatment(object):
class _MatchAllFn(beam.DoFn):
+
def __init__(self, empty_match_treatment):
self._empty_match_treatment = empty_match_treatment
@@ -183,6 +184,7 @@ class MatchFiles(beam.PTransform):
This ``PTransform`` returns a ``PCollection`` of matching files in the form
of ``FileMetadata`` objects."""
+
def __init__(
self,
file_pattern: str,
@@ -200,6 +202,7 @@ class MatchAll(beam.PTransform):
This ``PTransform`` returns a ``PCollection`` of matching files in the form
of ``FileMetadata`` objects."""
+
def __init__(self, empty_match_treatment=EmptyMatchTreatment.ALLOW):
self._empty_match_treatment = empty_match_treatment
@@ -212,6 +215,7 @@ class MatchAll(beam.PTransform):
class ReadableFile(object):
"""A utility class for accessing files."""
+
def __init__(self, metadata, compression=None):
self.metadata = metadata
self._compression = compression
@@ -231,6 +235,7 @@ class ReadableFile(object):
class _ReadMatchesFn(beam.DoFn):
+
def __init__(self, compression, skip_directories):
self._compression = compression
self._skip_directories = skip_directories
@@ -271,6 +276,7 @@ class MatchContinuously(beam.PTransform):
(https://cloud.google.com/storage/docs/pubsub-notifications)
when using GCS if possible.
"""
+
def __init__(
self,
file_pattern,
@@ -345,6 +351,7 @@ class ReadMatches(beam.PTransform):
"""Converts each result of MatchFiles() or MatchAll() to a ReadableFile.
This helps read in a file's contents or obtain a file descriptor."""
+
def __init__(self, compression=None, skip_directories=True):
self._compression = compression
self._skip_directories = skip_directories
@@ -373,6 +380,7 @@ class FileSink(object):
- The ``create_metadata`` method, which creates all metadata passed to
Filesystems.create.
"""
+
def create_metadata(
self, destination: str, full_file_name: str) -> FileMetadata:
return FileMetadata(
@@ -396,6 +404,7 @@ class TextSink(FileSink):
This sink simply calls file_handler.write(record.encode('utf8') + '\n') on all
records that come into it.
"""
+
def open(self, fh):
self._fh = fh
@@ -466,6 +475,7 @@ FileNaming = Callable[[Any, Any, int, int, Any, str, str], str]
def destination_prefix_naming(suffix=None) -> FileNaming:
+
def _inner(window, pane, shard_index, total_shards, compression, destination):
prefix = str(destination)
return _format_shard(
@@ -475,6 +485,7 @@ def destination_prefix_naming(suffix=None) -> FileNaming:
def default_file_naming(prefix, suffix=None) -> FileNaming:
+
def _inner(window, pane, shard_index, total_shards, compression, destination):
return _format_shard(
window, pane, shard_index, total_shards, compression, prefix, suffix)
@@ -483,6 +494,7 @@ def default_file_naming(prefix, suffix=None) -> FileNaming:
def single_file_naming(prefix, suffix=None) -> FileNaming:
+
def _inner(window, pane, shard_index, total_shards, compression, destination):
assert shard_index in (0, None), shard_index
assert total_shards in (1, None), total_shards
@@ -669,6 +681,7 @@ def _create_writer(
class _MoveTempFilesIntoFinalDestinationFn(beam.DoFn):
+
def __init__(self, path, file_naming_fn, temp_dir):
self.path = path
self.file_naming_fn = file_naming_fn
@@ -751,6 +764,7 @@ class _MoveTempFilesIntoFinalDestinationFn(beam.DoFn):
class _WriteShardedRecordsFn(beam.DoFn):
+
def __init__(
self, base_path, sink_fn: Callable[[Any], FileSink], shards: int):
self.base_path = base_path
@@ -795,6 +809,7 @@ class _WriteShardedRecordsFn(beam.DoFn):
class _AppendShardedDestination(beam.DoFn):
+
def __init__(self, destination: Callable[[Any], str], shards: int):
self.destination_fn = destination
self.shards = shards
@@ -879,12 +894,13 @@ class _WriteUnshardedRecordsFn(beam.DoFn):
sink.flush()
writer.close()
- file_result = FileResult(self._file_names[key],
- shard_index=-1,
- total_shards=0,
- window=key[1],
- pane=None, # TODO(pabloem): get the pane info
- destination=key[0])
+ file_result = FileResult(
+ self._file_names[key],
+ shard_index=-1,
+ total_shards=0,
+ window=key[1],
+ pane=None, # TODO(pabloem): get the pane info
+ destination=key[0])
yield beam.pvalue.TaggedOutput(
self.WRITTEN_FILES,
diff --git a/sdks/python/apache_beam/io/fileio_test.py b/sdks/python/apache_beam/io/fileio_test.py
index ff4be9d3d7c..4fc088746b4 100644
--- a/sdks/python/apache_beam/io/fileio_test.py
+++ b/sdks/python/apache_beam/io/fileio_test.py
@@ -59,6 +59,7 @@ def _get_file_reader(readable_file):
class MatchTest(_TestCaseWithTempDirCleanUp):
+
def test_basic_two_files(self):
files = []
tempdir = '%s%s' % (self._new_tempdir(), os.sep)
@@ -137,6 +138,7 @@ class MatchTest(_TestCaseWithTempDirCleanUp):
class ReadTest(_TestCaseWithTempDirCleanUp):
+
def test_basic_file_name_provided(self):
content = 'TestingMyContent\nIn multiple lines\nhaha!'
dir = '%s%s' % (self._new_tempdir(), os.sep)
@@ -322,6 +324,7 @@ class MatchIntegrationTest(unittest.TestCase):
class MatchContinuouslyTest(_TestCaseWithTempDirCleanUp):
+
def test_with_deduplication(self):
files = []
tempdir = '%s%s' % (self._new_tempdir(), os.sep)
@@ -461,6 +464,7 @@ class WriteFilesTest(_TestCaseWithTempDirCleanUp):
for elm in SIMPLE_COLLECTION}
class CsvSink(fileio.TextSink):
+
def __init__(self, headers):
self.headers = headers
@@ -469,6 +473,7 @@ class WriteFilesTest(_TestCaseWithTempDirCleanUp):
self._fh.write('\n'.encode('utf8'))
class JsonSink(fileio.TextSink):
+
def write(self, record):
self._fh.write(json.dumps(record).encode('utf8'))
self._fh.write('\n'.encode('utf8'))
@@ -497,8 +502,8 @@ class WriteFilesTest(_TestCaseWithTempDirCleanUp):
def test_write_to_dynamic_destination(self):
sink_params = [
- fileio.TextSink, # pass a type signature
- fileio.TextSink() # pass a FileSink object
+ fileio.TextSink, # pass a type signature
+ fileio.TextSink() # pass a FileSink object
]
for sink in sink_params:
@@ -522,8 +527,8 @@ class WriteFilesTest(_TestCaseWithTempDirCleanUp):
| fileio.ReadMatches()
| beam.Map(
lambda f: (
- os.path.basename(f.metadata.path).split('-')[0],
- sorted(map(int, f.read_utf8().strip().split('\n'))))))
+ os.path.basename(f.metadata.path).split('-')[0], sorted(
+ map(int, f.read_utf8().strip().split('\n'))))))
assert_that(
result,
@@ -664,7 +669,9 @@ class WriteFilesTest(_TestCaseWithTempDirCleanUp):
label='verifyApache')
def record_dofn(self):
+
class RecordDoFn(beam.DoFn):
+
def process(self, element):
WriteFilesTest.all_records.append(element)
diff --git a/sdks/python/apache_beam/io/filesystem.py b/sdks/python/apache_beam/io/filesystem.py
index bdc25dcf0fe..e7c649b4743 100644
--- a/sdks/python/apache_beam/io/filesystem.py
+++ b/sdks/python/apache_beam/io/filesystem.py
@@ -450,6 +450,7 @@ class FileMetadata(object):
last_updated_in_seconds: [Optional] last modified timestamp of the file, or
valued 0.0 if not specified.
"""
+
def __init__(
self,
path: str,
@@ -486,12 +487,14 @@ class MatchResult(object):
"""Result from the ``FileSystem`` match operation which contains the list
of matched ``FileMetadata``.
"""
+
def __init__(self, pattern: str, metadata_list: List[FileMetadata]) -> None:
self.metadata_list = metadata_list
self.pattern = pattern
class BeamIOError(IOError):
+
def __init__(self, msg, exception_details=None):
"""Class representing the errors thrown in the batch file operations.
Args:
diff --git a/sdks/python/apache_beam/io/filesystem_test.py b/sdks/python/apache_beam/io/filesystem_test.py
index a4d456a366d..731d276e164 100644
--- a/sdks/python/apache_beam/io/filesystem_test.py
+++ b/sdks/python/apache_beam/io/filesystem_test.py
@@ -42,6 +42,7 @@ from apache_beam.io.filesystem import FileSystem
class TestingFileSystem(FileSystem):
+
def __init__(self, pipeline_options, has_dirs=False):
super().__init__(pipeline_options)
self._has_dirs = has_dirs
@@ -112,6 +113,7 @@ class TestingFileSystem(FileSystem):
class TestFileSystem(unittest.TestCase):
+
def setUp(self):
self.fs = TestingFileSystem(pipeline_options=None)
@@ -210,8 +212,8 @@ class TestFileSystem(unittest.TestCase):
# It's a filter function of type (str, int) -> bool
# that returns true for expected objects
filter_func = expected_object_names
- expected_object_names = [(short_path, size) for short_path,
- size in objects
+ expected_object_names = [(short_path, size)
+ for short_path, size in objects
if filter_func(short_path, size)]
for object_name, size in objects:
@@ -219,8 +221,7 @@ class TestFileSystem(unittest.TestCase):
self.fs._insert_random_file(file_name, size)
expected_file_names = [('gs://%s/%s' % (bucket_name, object_name), size)
- for object_name,
- size in expected_object_names]
+ for object_name, size in expected_object_names]
actual_file_names = [
(file_metadata.path, file_metadata.size_in_bytes)
for file_metadata in self._flatten_match(self.fs.match([file_pattern]))
@@ -258,6 +259,7 @@ class TestFileSystem(unittest.TestCase):
class TestFileSystemWithDirs(TestFileSystem):
+
def setUp(self):
self.fs = TestingFileSystem(pipeline_options=None, has_dirs=True)
diff --git a/sdks/python/apache_beam/io/filesystemio.py b/sdks/python/apache_beam/io/filesystemio.py
index 571d1f2d269..340b5065d53 100644
--- a/sdks/python/apache_beam/io/filesystemio.py
+++ b/sdks/python/apache_beam/io/filesystemio.py
@@ -37,6 +37,7 @@ class Downloader(metaclass=abc.ABCMeta):
Implementations should support random access reads.
"""
+
@property
@abc.abstractmethod
def size(self):
@@ -60,6 +61,7 @@ class Downloader(metaclass=abc.ABCMeta):
class Uploader(metaclass=abc.ABCMeta):
"""Upload interface for a single file."""
+
@abc.abstractmethod
def put(self, data):
"""Write data to file sequentially.
@@ -81,6 +83,7 @@ class Uploader(metaclass=abc.ABCMeta):
class DownloaderStream(io.RawIOBase):
"""Provides a stream interface for Downloader objects."""
+
def __init__(
self, downloader, read_buffer_size=io.DEFAULT_BUFFER_SIZE, mode='rb'):
"""Initializes the stream.
@@ -174,6 +177,7 @@ class DownloaderStream(io.RawIOBase):
class UploaderStream(io.RawIOBase):
"""Provides a stream interface for Uploader objects."""
+
def __init__(self, uploader, mode='wb'):
"""Initializes the stream.
@@ -228,6 +232,7 @@ class PipeStream(object):
Remembers the last ``size`` bytes read and allows rewinding the stream by that
amount exactly. See BEAM-6380 for more.
"""
+
def __init__(self, recv_pipe):
self.conn = recv_pipe
self.closed = False
diff --git a/sdks/python/apache_beam/io/filesystemio_test.py b/sdks/python/apache_beam/io/filesystemio_test.py
index cbf9449e78d..48206687c0b 100644
--- a/sdks/python/apache_beam/io/filesystemio_test.py
+++ b/sdks/python/apache_beam/io/filesystemio_test.py
@@ -32,6 +32,7 @@ _LOGGER = logging.getLogger(__name__)
class FakeDownloader(filesystemio.Downloader):
+
def __init__(self, data):
self._data = data
self.last_read_size = -1
@@ -46,6 +47,7 @@ class FakeDownloader(filesystemio.Downloader):
class FakeUploader(filesystemio.Uploader):
+
def __init__(self):
self.data = b''
self.last_write_size = -1
@@ -64,6 +66,7 @@ class FakeUploader(filesystemio.Uploader):
class TestDownloaderStream(unittest.TestCase):
+
def test_file_attributes(self):
downloader = FakeDownloader(data=None)
stream = filesystemio.DownloaderStream(downloader)
@@ -102,6 +105,7 @@ class TestDownloaderStream(unittest.TestCase):
class TestUploaderStream(unittest.TestCase):
+
def test_file_attributes(self):
uploader = FakeUploader()
stream = filesystemio.UploaderStream(uploader)
@@ -147,6 +151,7 @@ class TestUploaderStream(unittest.TestCase):
class TestPipeStream(unittest.TestCase):
+
def _read_and_verify(self, stream, expected, buffer_size, success):
data_list = []
bytes_read = 0
diff --git a/sdks/python/apache_beam/io/filesystems_test.py b/sdks/python/apache_beam/io/filesystems_test.py
index 1ea7c34d9f4..66bf316c722 100644
--- a/sdks/python/apache_beam/io/filesystems_test.py
+++ b/sdks/python/apache_beam/io/filesystems_test.py
@@ -36,6 +36,7 @@ from apache_beam.io.filesystems import FileSystems
def _gen_fake_join(separator):
"""Returns a callable that joins paths with the given separator."""
+
def _join(first_path, *paths):
return separator.join((first_path.rstrip(separator), ) + paths)
@@ -43,6 +44,7 @@ def _gen_fake_join(separator):
class FileSystemsTest(unittest.TestCase):
+
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
diff --git a/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source_test.py b/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source_test.py
index f6ba5abdd57..15405e459e4 100644
--- a/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source_test.py
+++ b/sdks/python/apache_beam/io/flink/flink_streaming_impulse_source_test.py
@@ -26,6 +26,7 @@ from apache_beam.io.flink.flink_streaming_impulse_source import FlinkStreamingIm
class FlinkStreamingImpulseSourceTest(unittest.TestCase):
+
def test_serialization(self):
p = beam.Pipeline()
# pylint: disable=expression-not-assigned
diff --git a/sdks/python/apache_beam/io/gcp/__init__.py b/sdks/python/apache_beam/io/gcp/__init__.py
index f88a0117aa4..bf5dfd0acb5 100644
--- a/sdks/python/apache_beam/io/gcp/__init__.py
+++ b/sdks/python/apache_beam/io/gcp/__init__.py
@@ -25,7 +25,9 @@ try:
from apitools.base.py import transfer
class _WrapperNamespace(object):
+
class BytesGenerator(email_generator.BytesGenerator):
+
def _write_lines(self, lines):
self.write(lines)
diff --git a/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py b/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py
index 052790c4a20..1b85997dd72 100644
--- a/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py
@@ -89,6 +89,7 @@ DIALECT_OUTPUT_EXPECTED = [('apple', ), ('orange', )]
class BigQueryQueryToTableIT(unittest.TestCase):
+
def setUp(self):
self.test_pipeline = TestPipeline(is_integration_test=True)
self.runner_name = type(self.test_pipeline.runner).__name__
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py
index 9f60b5af672..ce8441e2982 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery.py
@@ -520,6 +520,7 @@ class TableRowJsonCoder(coders.Coder):
table schema in order to obtain the ordered list of field names. Reading from
sources on the other hand does not need the table schema.
"""
+
def __init__(self, table_schema=None):
# The table schema is needed for encoding TableRows as JSON (writing to
# sinks) because the ordered list of field names is used in the JSON
@@ -642,6 +643,7 @@ class _BigQueryExportResult:
class _CustomBigQuerySource(BoundedSource):
+
def __init__(
self,
method,
@@ -829,8 +831,10 @@ class _CustomBigQuerySource(BoundedSource):
weight=1.0, source=source, start_position=None, stop_position=None)
def get_range_tracker(self, start_position, stop_position):
+
class CustomBigQuerySourceRangeTracker(RangeTracker):
"""A RangeTracker that always returns positions as None."""
+
def start_position(self):
return None
@@ -1221,9 +1225,11 @@ class _CustomBigQueryStorageSource(BoundedSource):
weight=1.0, source=source, start_position=None, stop_position=None)
def get_range_tracker(self, start_position, stop_position):
+
class NonePositionRangeTracker(RangeTracker):
"""A RangeTracker that always returns positions as None. Prevents the
BigQuery Storage source from being read() before being split()."""
+
def start_position(self):
return None
@@ -1341,6 +1347,7 @@ class _CustomBigQueryStorageStreamSource(BoundedSource):
class _ReadReadRowsResponsesWithFastAvro():
"""An iterator that deserializes ReadRowsResponses using the fastavro
library."""
+
def __init__(self, read_rows_iterator, read_rows_response):
self.read_rows_iterator = read_rows_iterator
self.read_rows_response = read_rows_response
@@ -1617,8 +1624,8 @@ class BigQueryWriteFn(DoFn):
# Flush current batch first if adding this row will exceed our limits
# limits: byte size; number of rows
- if ((self._destination_buffer_byte_size[destination] + row_byte_size >
- self._max_insert_payload_size) or
+ if ((self._destination_buffer_byte_size[destination] + row_byte_size
+ > self._max_insert_payload_size) or
len(self._rows_buffer[destination]) >= self._max_batch_size):
flushed_batch = self._flush_batch(destination)
# After flushing our existing batch, we now buffer the current row
@@ -1712,9 +1719,8 @@ class BigQueryWriteFn(DoFn):
# - WARNING when we are continuing to retry, and have a deadline.
# - ERROR when we will no longer retry, or MAY retry forever.
log_level = (
- logging.WARN if should_retry or
- self._retry_strategy != RetryStrategy.RETRY_ALWAYS else
- logging.ERROR)
+ logging.WARN if should_retry or self._retry_strategy
+ != RetryStrategy.RETRY_ALWAYS else logging.ERROR)
_LOGGER.log(log_level, message)
@@ -1740,16 +1746,13 @@ class BigQueryWriteFn(DoFn):
[
pvalue.TaggedOutput(
BigQueryWriteFn.FAILED_ROWS_WITH_ERRORS,
- w.with_value((destination, row, err))) for row,
- err,
- w in failed_rows
+ w.with_value((destination, row, err)))
+ for row, err, w in failed_rows
],
[
pvalue.TaggedOutput(
BigQueryWriteFn.FAILED_ROWS, w.with_value((destination, row)))
- for row,
- unused_err,
- w in failed_rows
+ for row, unused_err, w in failed_rows
])
@@ -1761,6 +1764,7 @@ DEFAULT_BATCH_BUFFERING_DURATION_LIMIT_SEC = 0.2
class _StreamToBigQuery(PTransform):
+
def __init__(
self,
table_reference,
@@ -1801,6 +1805,7 @@ class _StreamToBigQuery(PTransform):
self._max_insert_payload_size = max_insert_payload_size
class InsertIdPrefixFn(DoFn):
+
def start_bundle(self):
self.prefix = str(uuid.uuid4())
self._row_count = 0
@@ -1891,6 +1896,7 @@ class WriteToBigQuery(PTransform):
tables. The elements would come in as Python dictionaries, or as `TableRow`
instances.
"""
+
class Method(object):
DEFAULT = 'DEFAULT'
STREAMING_INSERTS = 'STREAMING_INSERTS'
@@ -2331,10 +2337,9 @@ bigquery_v2_messages.TableSchema`. or a `ValueProvider` that has a JSON string,
# remove_objects_from_args and insert_values_in_args
# are currently implemented.
def serialize(side_inputs):
- return {(SIDE_INPUT_PREFIX + '%s') % ix:
- si.to_runner_api(context).SerializeToString()
- for ix,
- si in enumerate(side_inputs)}
+ return {(SIDE_INPUT_PREFIX + '%s') % ix: si.to_runner_api(
+ context).SerializeToString()
+ for ix, si in enumerate(side_inputs)}
table_side_inputs = serialize(self.table_side_inputs)
schema_side_inputs = serialize(self.schema_side_inputs)
@@ -2382,8 +2387,8 @@ bigquery_v2_messages.TableSchema`. or a `ValueProvider` that has a JSON string,
# to_runner_api_parameter above).
indexed_side_inputs = [(
get_sideinput_index(tag),
- pvalue.AsSideInput.from_runner_api(si, context)) for tag,
- si in deserialized_side_inputs.items()]
+ pvalue.AsSideInput.from_runner_api(si, context))
+ for tag, si in deserialized_side_inputs.items()]
return [si for _, si in sorted(indexed_side_inputs)]
config['table_side_inputs'] = deserialize(config['table_side_inputs'])
@@ -2395,6 +2400,7 @@ bigquery_v2_messages.TableSchema`. or a `ValueProvider` that has a JSON string,
class WriteResult:
"""The result of a WriteToBigQuery transform.
"""
+
def __init__(
self,
method: str = None,
@@ -2661,8 +2667,8 @@ class StorageWriteToBigQuery(PTransform):
failed_rows = failed_rows | beam.Map(lambda row: row.as_dict())
failed_rows_with_errors = failed_rows_with_errors | beam.Map(
lambda row: {
- "error_message": row.error_message,
- "failed_row": row.failed_row.as_dict()
+ "error_message": row.error_message, "failed_row": row.failed_row.
+ as_dict()
})
return WriteResult(
@@ -2671,6 +2677,7 @@ class StorageWriteToBigQuery(PTransform):
failed_rows_with_errors=failed_rows_with_errors)
class ConvertToBeamRows(PTransform):
+
def __init__(self, schema, dynamic_destinations):
self.schema = schema
self.dynamic_destinations = dynamic_destinations
@@ -2682,8 +2689,8 @@ class StorageWriteToBigQuery(PTransform):
| "Convert dict to Beam Row" >> beam.Map(
lambda row: beam.Row(
**{
- StorageWriteToBigQuery.DESTINATION: row[0],
- StorageWriteToBigQuery.RECORD: bigquery_tools.
+ StorageWriteToBigQuery.DESTINATION: row[
+ 0], StorageWriteToBigQuery.RECORD: bigquery_tools.
beam_row_from_dict(row[1], self.schema)
})))
else:
@@ -2804,6 +2811,7 @@ class ReadFromBigQuery(PTransform):
`BEAM_ROW`. For more information on schemas, see
https://beam.apache.org/documentation/programming-guide/#what-is-a-schema)
"""
+
class Method(object):
EXPORT = 'EXPORT' # This is currently the default.
DIRECT_READ = 'DIRECT_READ'
@@ -2972,6 +2980,7 @@ class ReadFromBigQueryRequest:
"""
Class that defines data to read from BQ.
"""
+
def __init__(
self,
query: str = None,
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_avro_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_avro_tools_test.py
index eca208d2661..36f652ad7ce 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_avro_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_avro_tools_test.py
@@ -27,6 +27,7 @@ from apache_beam.io.gcp.internal.clients import bigquery
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestBigQueryToAvroSchema(unittest.TestCase):
+
def test_convert_bigquery_schema_to_avro_schema(self):
subfields = [
bigquery.TableFieldSchema(
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
index 3145fb51106..2a6156eb462 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py
@@ -92,6 +92,7 @@ def _generate_job_name(job_name, job_type, step_name):
def file_prefix_generator(
with_validation=True, pipeline_gcs_location=None, temp_location=None):
+
def _generate_file_prefix(unused_elm):
# If a gcs location is provided to the pipeline, then we shall use that.
# Otherwise, we shall use the temp_location from pipeline options.
@@ -291,6 +292,7 @@ class WriteGroupedRecordsToFile(beam.DoFn):
Experimental; no backwards compatibility guarantees.
"""
+
def __init__(
self, schema, max_file_size=_DEFAULT_MAX_FILE_SIZE, file_format=None):
self.schema = schema
@@ -344,6 +346,7 @@ class UpdateDestinationSchema(beam.DoFn):
Experimental; no backwards compatibility guarantees.
"""
+
def __init__(
self,
project=None,
@@ -790,6 +793,7 @@ class PartitionFiles(beam.DoFn):
SINGLE_PARTITION_TAG = 'SINGLE_PARTITION'
class Partition(object):
+
def __init__(self, max_size, max_files, files=None, size=0):
self.max_size = max_size
self.max_files = max_files
@@ -848,6 +852,7 @@ class PartitionFiles(beam.DoFn):
class DeleteTablesFn(beam.DoFn):
+
def __init__(self, test_client=None):
self.test_client = test_client
@@ -1145,8 +1150,7 @@ class BigQueryBatchFileLoads(beam.PTransform):
# https://github.com/apache/beam/issues/24535.
finished_temp_tables_load_job_ids_list_pc = (
finished_temp_tables_load_job_ids_pc | beam.MapTuple(
- lambda destination,
- job_reference: (
+ lambda destination, job_reference: (
bigquery_tools.parse_table_reference(destination).tableId,
(destination, job_reference)))
| beam.GroupByKey()
@@ -1234,8 +1238,7 @@ class BigQueryBatchFileLoads(beam.PTransform):
singleton_pc
| "SchemaModJobNamePrefix" >> beam.Map(
lambda _: _generate_job_name(
- job_name,
- bigquery_tools.BigQueryJobTypes.LOAD,
+ job_name, bigquery_tools.BigQueryJobTypes.LOAD,
'SCHEMA_MOD_STEP')))
copy_job_name_pcv = pvalue.AsSingleton(
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
index 10453d9c8ba..051aee5507e 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py
@@ -254,6 +254,7 @@ class TestWriteRecordsToFile(_TestCaseWithTempDirCleanUp):
class TestWriteGroupedRecordsToFile(_TestCaseWithTempDirCleanUp):
+
def _consume_input(self, fn, input, checks):
if checks is None:
return
@@ -404,6 +405,7 @@ class TestPartitionFiles(unittest.TestCase):
class TestBigQueryFileLoads(_TestCaseWithTempDirCleanUp):
+
def test_trigger_load_jobs_with_empty_files(self):
destination = "project:dataset.table"
empty_files = []
@@ -799,6 +801,7 @@ class TestBigQueryFileLoads(_TestCaseWithTempDirCleanUp):
# Insert a fake clock to work with auto-sharding which needs a processing
# time timer.
class _FakeClock(object):
+
def __init__(self, now=time.time()):
self._now = now
@@ -827,12 +830,12 @@ class TestBigQueryFileLoads(_TestCaseWithTempDirCleanUp):
if is_streaming:
_SIZE = len(_ELEMENTS)
fisrt_batch = [
- TimestampedValue(value, start_time + i + 1) for i,
- value in enumerate(_ELEMENTS[:_SIZE // 2])
+ TimestampedValue(value, start_time + i + 1)
+ for i, value in enumerate(_ELEMENTS[:_SIZE // 2])
]
second_batch = [
- TimestampedValue(value, start_time + _SIZE // 2 + i + 1) for i,
- value in enumerate(_ELEMENTS[_SIZE // 2:])
+ TimestampedValue(value, start_time + _SIZE // 2 + i + 1)
+ for i, value in enumerate(_ELEMENTS[_SIZE // 2:])
]
# Advance processing time between batches of input elements to fire the
# user triggers. Intentionally advance the processing time twice for the
@@ -1031,12 +1034,10 @@ class BigQueryFileLoadsIT(unittest.TestCase):
_ = (
input | "WriteWithMultipleDestsFreely" >> bigquery.WriteToBigQuery(
- table=lambda x,
- tables:
+ table=lambda x, tables:
(tables['table1'] if 'language' in x else tables['table2']),
table_side_inputs=(table_record_pcv, ),
- schema=lambda dest,
- schema_map: schema_map.get(dest, None),
+ schema=lambda dest, schema_map: schema_map.get(dest, None),
schema_side_inputs=(schema_map_pcv, ),
create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY))
@@ -1045,8 +1046,7 @@ class BigQueryFileLoadsIT(unittest.TestCase):
input | "WriteWithMultipleDests" >> bigquery.WriteToBigQuery(
table=lambda x:
(output_table_3 if 'language' in x else output_table_4),
- schema=lambda dest,
- schema_map: schema_map.get(dest, None),
+ schema=lambda dest, schema_map: schema_map.get(dest, None),
schema_side_inputs=(schema_map_pcv, ),
create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
write_disposition=beam.io.BigQueryDisposition.WRITE_EMPTY,
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_io_metadata.py b/sdks/python/apache_beam/io/gcp/bigquery_io_metadata.py
index a730f2cfc9b..c98a841c3a6 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_io_metadata.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_io_metadata.py
@@ -80,6 +80,7 @@ class BigQueryIOMetadata(object):
Do not construct directly, use the create_bigquery_io_metadata factory.
Which will request metadata properly based on which runner is being used.
"""
+
def __init__(self, beam_job_id=None, step_name=None):
self.beam_job_id = beam_job_id
self.step_name = step_name
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_io_metadata_test.py b/sdks/python/apache_beam/io/gcp/bigquery_io_metadata_test.py
index c9120246538..a83b2e0a78b 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_io_metadata_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_io_metadata_test.py
@@ -26,6 +26,7 @@ from apache_beam.io.gcp import bigquery_io_metadata
class BigqueryIoMetadataTest(unittest.TestCase):
+
def test_is_valid_cloud_label_value(self):
# A dataflow job ID.
# Lowercase letters, numbers, underscores and hyphens are allowed.
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py b/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py
index 8ca9736b025..b294c540125 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_io_read_pipeline.py
@@ -37,6 +37,7 @@ from apache_beam.testing.util import equal_to
class RowToStringWithSlowDown(beam.DoFn):
+
def process(self, element, num_slow=0, *args, **kwargs):
if num_slow == 0:
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_json_it_test.py b/sdks/python/apache_beam/io/gcp/bigquery_json_it_test.py
index 6716aa1bc10..e1e71aa875d 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_json_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_json_it_test.py
@@ -137,6 +137,7 @@ class BigQueryJsonIT(unittest.TestCase):
return json.loads(value)
class CompareJson(beam.DoFn, unittest.TestCase):
+
def process(self, row):
country_code = row["country_code"]
expected = json_data[country_code]
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
index f038b48e04d..0ab27650a9f 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py
@@ -109,15 +109,19 @@ class _PassThroughThenCleanup(PTransform):
Utilizes readiness of PCollection to trigger DoFn.
"""
+
def __init__(self, side_input=None):
self.side_input = side_input
def expand(self, input):
+
class PassThrough(beam.DoFn):
+
def process(self, element):
yield element
class RemoveExtractedFiles(beam.DoFn):
+
def process(self, unused_element, unused_signal, gcs_locations):
FileSystems.delete(list(gcs_locations))
@@ -144,6 +148,7 @@ class _PassThroughThenCleanupTempDatasets(PTransform):
Utilizes readiness of PCollection to trigger DoFn.
"""
+
def __init__(self, side_input=None):
self.side_input = side_input
@@ -151,10 +156,12 @@ class _PassThroughThenCleanupTempDatasets(PTransform):
pipeline_options = input.pipeline.options
class PassThrough(beam.DoFn):
+
def process(self, element):
yield element
class CleanUpProjects(beam.DoFn):
+
def process(self, unused_element, unused_signal, pipeline_details):
bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(
pipeline_options)
@@ -191,6 +198,7 @@ class _BigQueryReadSplit(beam.transforms.DoFn):
This transform will start a BigQuery export job, and output a number of
file sources that are consumed downstream.
"""
+
def __init__(
self,
options: PipelineOptions,
@@ -404,6 +412,7 @@ FieldSchema = collections.namedtuple('FieldSchema', 'fields mode name type')
class _JsonToDictCoder(coders.Coder):
"""A coder for a JSON string to a Python dict."""
+
def __init__(self, table_schema):
self.fields = self._convert_to_tuple(table_schema.fields)
self._converters = {
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py b/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py
index 913d6e078d8..7fc6ab58142 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_it_test.py
@@ -62,6 +62,7 @@ def skip(runners):
runners = [runners]
def inner(fn):
+
@wraps(fn)
def wrapped(self):
if self.runner_name in runners:
@@ -603,6 +604,7 @@ class ReadUsingStorageApiTests(BigQueryReadIntegrationTests):
class ReadNewTypesTests(BigQueryReadIntegrationTests):
+
@classmethod
def setUpClass(cls):
super(ReadNewTypesTests, cls).setUpClass()
@@ -824,6 +826,7 @@ class ReadAllBQTests(BigQueryReadIntegrationTests):
class ReadInteractiveRunnerTests(BigQueryReadIntegrationTests):
+
@skip(['PortableRunner', 'FlinkRunner'])
@pytest.mark.it_postcommit
def test_read_in_interactive_runner(self):
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py b/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py
index 0b4cfe2ecba..2f989c7c368 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_read_perf_test.py
@@ -79,6 +79,7 @@ except ImportError:
class BigQueryReadPerfTest(LoadTest):
+
def __init__(self):
super().__init__()
self.input_dataset = self.pipeline.get_option('input_dataset')
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py
index beb373a7dea..7d3acd4d16c 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py
@@ -101,6 +101,7 @@ def convert_to_usertype(table_schema, selected_fields=None):
class BeamSchemaConversionDoFn(DoFn):
+
def __init__(self, pcoll_val_ctor):
self._pcoll_val_ctor = pcoll_val_ctor
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py
index 7ae49dff205..7079d5b4f15 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools_test.py
@@ -35,6 +35,7 @@ except ImportError:
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestBigQueryToSchema(unittest.TestCase):
+
def test_check_schema_conversions(self):
fields = [
bigquery.TableFieldSchema(name='stn', type='STRING', mode="NULLABLE"),
@@ -173,6 +174,7 @@ class TestBigQueryToSchema(unittest.TestCase):
table=value_provider.ValueProvider(), output_type='BEAM_ROW')
def test_unsupported_callable(self):
+
def filterTable(table):
if table is not None:
return table
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py
index 435fe67d02f..8e5070126d5 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py
@@ -148,6 +148,7 @@ def _load_or_default(filename):
HttpError is None or gcp_bigquery is None,
'GCP dependencies are not installed')
class TestTableRowJsonCoder(unittest.TestCase):
+
def test_row_as_table_row(self):
schema_definition = [('s', 'STRING'), ('i', 'INTEGER'), ('f', 'FLOAT'),
('b', 'BOOLEAN'), ('n', 'NUMERIC'), ('r', 'RECORD'),
@@ -172,8 +173,8 @@ class TestTableRowJsonCoder(unittest.TestCase):
'"g": "LINESTRING(1 2, 3 4, 5 6, 7 8)"}')
schema = bigquery.TableSchema(
fields=[
- bigquery.TableFieldSchema(name=k, type=v) for k,
- v in schema_definition
+ bigquery.TableFieldSchema(name=k, type=v)
+ for k, v in schema_definition
])
coder = TableRowJsonCoder(table_schema=schema)
@@ -211,8 +212,8 @@ class TestTableRowJsonCoder(unittest.TestCase):
schema_definition = [('f', 'FLOAT')]
schema = bigquery.TableSchema(
fields=[
- bigquery.TableFieldSchema(name=k, type=v) for k,
- v in schema_definition
+ bigquery.TableFieldSchema(name=k, type=v)
+ for k, v in schema_definition
])
coder = TableRowJsonCoder(table_schema=schema)
test_row = bigquery.TableRow(
@@ -231,8 +232,10 @@ class TestTableRowJsonCoder(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestJsonToDictCoder(unittest.TestCase):
+
@staticmethod
def _make_schema(fields):
+
def _fill_schema(fields):
for field in fields:
table_field = bigquery.TableFieldSchema()
@@ -334,9 +337,12 @@ class TestJsonToDictCoder(unittest.TestCase):
HttpError is None or HttpForbiddenError is None,
'GCP dependencies are not installed')
class TestReadFromBigQuery(unittest.TestCase):
+
@classmethod
def setUpClass(cls):
+
class UserDefinedOptions(PipelineOptions):
+
@classmethod
def _add_argparse_args(cls, parser):
parser.add_value_provider_argument('--gcs_location')
@@ -507,7 +513,9 @@ class TestReadFromBigQuery(unittest.TestCase):
expected_retries=3),
])
def test_get_table_transient_exception(self, responses, expected_retries):
+
class DummyTable:
+
class DummySchema:
fields = []
@@ -619,7 +627,9 @@ class TestReadFromBigQuery(unittest.TestCase):
expected_retries=2),
])
def test_get_table_non_transient_exception(self, responses, expected_retries):
+
class DummyTable:
+
class DummySchema:
fields = []
@@ -765,6 +775,7 @@ class TestReadFromBigQuery(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestBigQuerySink(unittest.TestCase):
+
def test_table_spec_display_data(self):
sink = beam.io.BigQuerySink('dataset.table')
dd = DisplayData.create_from(sink)
@@ -796,6 +807,7 @@ class TestBigQuerySink(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestWriteToBigQuery(unittest.TestCase):
+
def _cleanup_files(self):
if os.path.exists('insert_calls1'):
os.remove('insert_calls1')
@@ -962,8 +974,7 @@ class TestWriteToBigQuery(unittest.TestCase):
schema = value_provider.StaticValueProvider(str, '"a:str"')
original = WriteToBigQuery(
- table=lambda _,
- side_input: side_input['table'],
+ table=lambda _, side_input: side_input['table'],
table_side_inputs=(table_record_pcv, ),
schema=schema)
@@ -978,8 +989,7 @@ class TestWriteToBigQuery(unittest.TestCase):
# Find the transform from the context.
write_to_bq_id = [
- k for k,
- v in pipeline_proto.components.transforms.items()
+ k for k, v in pipeline_proto.components.transforms.items()
if v.unique_name == 'MyWriteToBigQuery'
][0]
deserialized_node = context.transforms.get_by_id(write_to_bq_id)
@@ -1010,6 +1020,7 @@ class TestWriteToBigQuery(unittest.TestCase):
original_side_input_data.view_fn, deserialized_side_input_data.view_fn)
def test_streaming_triggering_frequency_without_auto_sharding(self):
+
def noop(table, **kwargs):
return []
@@ -1039,6 +1050,7 @@ class TestWriteToBigQuery(unittest.TestCase):
test_client=client))
def test_streaming_triggering_frequency_with_auto_sharding(self):
+
def noop(table, **kwargs):
return []
@@ -1219,25 +1231,27 @@ class BigQueryStreamingInsertsErrorHandling(unittest.TestCase):
# failed rows
param(
insert_response=[
- exceptions.TooManyRequests if exceptions else None,
- None],
- error_reason='Too Many Requests', # not in _NON_TRANSIENT_ERRORS
+ exceptions.TooManyRequests if exceptions else None, None
+ ],
+ error_reason='Too Many Requests', # not in _NON_TRANSIENT_ERRORS
failed_rows=[]),
# reason not in _NON_TRANSIENT_ERRORS for row 1 on both attempts, sent to
# failed rows after hitting max_retries
param(
insert_response=[
- exceptions.InternalServerError if exceptions else None,
- exceptions.InternalServerError if exceptions else None],
- error_reason='Internal Server Error', # not in _NON_TRANSIENT_ERRORS
+ exceptions.InternalServerError if exceptions else None,
+ exceptions.InternalServerError if exceptions else None
+ ],
+ error_reason='Internal Server Error', # not in _NON_TRANSIENT_ERRORS
failed_rows=['value1', 'value3', 'value5']),
# reason in _NON_TRANSIENT_ERRORS for row 1 on both attempts, sent to
# failed_rows after hitting max_retries
param(
insert_response=[
- exceptions.Forbidden if exceptions else None,
- exceptions.Forbidden if exceptions else None],
- error_reason='Forbidden', # in _NON_TRANSIENT_ERRORS
+ exceptions.Forbidden if exceptions else None,
+ exceptions.Forbidden if exceptions else None
+ ],
+ error_reason='Forbidden', # in _NON_TRANSIENT_ERRORS
failed_rows=['value1', 'value3', 'value5']),
])
def test_insert_rows_json_exception_retry_always(
@@ -1363,63 +1377,63 @@ class BigQueryStreamingInsertsErrorHandling(unittest.TestCase):
@parameterized.expand([
param(
exception_type=exceptions.DeadlineExceeded if exceptions else None,
- error_reason='Deadline Exceeded', # not in _NON_TRANSIENT_ERRORS
+ error_reason='Deadline Exceeded', # not in _NON_TRANSIENT_ERRORS
failed_values=[],
expected_call_count=2),
param(
exception_type=exceptions.Conflict if exceptions else None,
- error_reason='Conflict', # not in _NON_TRANSIENT_ERRORS
+ error_reason='Conflict', # not in _NON_TRANSIENT_ERRORS
failed_values=[],
expected_call_count=2),
param(
exception_type=exceptions.TooManyRequests if exceptions else None,
- error_reason='Too Many Requests', # not in _NON_TRANSIENT_ERRORS
+ error_reason='Too Many Requests', # not in _NON_TRANSIENT_ERRORS
failed_values=[],
expected_call_count=2),
param(
exception_type=exceptions.InternalServerError if exceptions else None,
- error_reason='Internal Server Error', # not in _NON_TRANSIENT_ERRORS
+ error_reason='Internal Server Error', # not in _NON_TRANSIENT_ERRORS
failed_values=[],
expected_call_count=2),
param(
exception_type=exceptions.BadGateway if exceptions else None,
- error_reason='Bad Gateway', # not in _NON_TRANSIENT_ERRORS
+ error_reason='Bad Gateway', # not in _NON_TRANSIENT_ERRORS
failed_values=[],
expected_call_count=2),
param(
exception_type=exceptions.ServiceUnavailable if exceptions else None,
- error_reason='Service Unavailable', # not in _NON_TRANSIENT_ERRORS
+ error_reason='Service Unavailable', # not in _NON_TRANSIENT_ERRORS
failed_values=[],
expected_call_count=2),
param(
exception_type=exceptions.GatewayTimeout if exceptions else None,
- error_reason='Gateway Timeout', # not in _NON_TRANSIENT_ERRORS
+ error_reason='Gateway Timeout', # not in _NON_TRANSIENT_ERRORS
failed_values=[],
expected_call_count=2),
param(
exception_type=exceptions.BadRequest if exceptions else None,
- error_reason='Bad Request', # in _NON_TRANSIENT_ERRORS
+ error_reason='Bad Request', # in _NON_TRANSIENT_ERRORS
failed_values=['value1', 'value2'],
expected_call_count=1),
param(
exception_type=exceptions.Unauthorized if exceptions else None,
- error_reason='Unauthorized', # in _NON_TRANSIENT_ERRORS
+ error_reason='Unauthorized', # in _NON_TRANSIENT_ERRORS
failed_values=['value1', 'value2'],
expected_call_count=1),
param(
exception_type=exceptions.Forbidden if exceptions else None,
- error_reason='Forbidden', # in _NON_TRANSIENT_ERRORS
+ error_reason='Forbidden', # in _NON_TRANSIENT_ERRORS
failed_values=['value1', 'value2'],
expected_call_count=1),
param(
exception_type=exceptions.NotFound if exceptions else None,
- error_reason='Not Found', # in _NON_TRANSIENT_ERRORS
+ error_reason='Not Found', # in _NON_TRANSIENT_ERRORS
failed_values=['value1', 'value2'],
expected_call_count=1),
param(
exception_type=exceptions.MethodNotImplemented
- if exceptions else None,
- error_reason='Not Implemented', # in _NON_TRANSIENT_ERRORS
+ if exceptions else None,
+ error_reason='Not Implemented', # in _NON_TRANSIENT_ERRORS
failed_values=['value1', 'value2'],
expected_call_count=1),
])
@@ -1915,6 +1929,7 @@ class BigQueryStreamingInsertsErrorHandling(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class BigQueryStreamingInsertTransformTests(unittest.TestCase):
+
def test_dofn_client_process_performs_batching(self):
client = mock.Mock()
client.tables.Get.return_value = bigquery.Table(
@@ -2051,6 +2066,7 @@ class BigQueryStreamingInsertTransformTests(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class PipelineBasedStreamingInsertTest(_TestCaseWithTempDirCleanUp):
+
@mock.patch('time.sleep')
def test_failure_has_same_insert_ids(self, unused_mock_sleep):
tempdir = '%s%s' % (self._new_tempdir(), os.sep)
@@ -2460,12 +2476,10 @@ class BigQueryStreamingInsertTransformIntegrationTests(unittest.TestCase):
r = (
input
| "WriteWithMultipleDests" >> beam.io.gcp.bigquery.WriteToBigQuery(
- table=lambda x,
- tables:
+ table=lambda x, tables:
(tables['table1'] if 'language' in x else tables['table2']),
table_side_inputs=(table_record_pcv, ),
- schema=lambda dest,
- table_map: table_map.get(dest, None),
+ schema=lambda dest, table_map: table_map.get(dest, None),
schema_side_inputs=(schema_table_pcv, ),
insert_retry_strategy=RetryStrategy.RETRY_ON_TRANSIENT_ERROR,
method='STREAMING_INSERTS'))
@@ -2665,8 +2679,7 @@ class BigQueryFileLoadsIntegrationTests(unittest.TestCase):
input
| 'WriteToBigQuery' >> beam.io.gcp.bigquery.WriteToBigQuery(
table='%s:%s' % (self.project, self.output_table),
- schema=lambda _,
- schema: schema,
+ schema=lambda _, schema: schema,
schema_side_inputs=(beam.pvalue.AsSingleton(schema_pc), ),
method='FILE_LOADS',
temp_file_format=bigquery_tools.FileFormat.AVRO,
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index 48da929a07b..4686ded89b5 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -1282,8 +1282,8 @@ class BigQueryWrapper(object):
# can happen during retries on failures.
# TODO(silviuc): Must add support to writing TableRow's instead of dicts.
insert_ids = [
- str(self.unique_row_id) if not insert_ids else insert_ids[i] for i,
- _ in enumerate(rows)
+ str(self.unique_row_id) if not insert_ids else insert_ids[i]
+ for i, _ in enumerate(rows)
]
rows = [
fast_json_loads(fast_json_dumps(r, default=default_encoder))
@@ -1392,6 +1392,7 @@ class RowAsDictJsonCoder(coders.Coder):
This is the default coder for sources and sinks if the coder argument is not
specified.
"""
+
def encode(self, table_row):
# The normal error when dumping NAN/INF values is:
# ValueError: Out of range float values are not JSON compliant
@@ -1419,6 +1420,7 @@ class JsonRowWriter(io.IOBase):
A writer which provides an IOBase-like interface for writing table rows
(represented as dicts) as newline-delimited JSON strings.
"""
+
def __init__(self, file_handle):
"""Initialize an JsonRowWriter.
@@ -1459,6 +1461,7 @@ class AvroRowWriter(io.IOBase):
A writer which provides an IOBase-like interface for writing table rows
(represented as dicts) as Avro records.
"""
+
def __init__(self, file_handle, schema):
"""Initialize an AvroRowWriter.
@@ -1552,6 +1555,7 @@ class AppendDestinationsFn(DoFn):
Experimental; no backwards compatibility guarantees.
"""
+
def __init__(self, destination):
self._display_destination = destination
self.destination = AppendDestinationsFn._get_table_fn(destination)
@@ -1656,6 +1660,7 @@ bigquery_v2_messages.TableSchema` format.
def table_schema_to_dict(table_schema):
"""Create a dictionary representation of table schema for serialization
"""
+
def get_table_field(field):
"""Create a dictionary representation of a table field
"""
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
index 1307a788692..c463186e6f7 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -72,6 +72,7 @@ except ImportError:
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestTableSchemaParser(unittest.TestCase):
+
def test_parse_table_schema_from_json(self):
string_field = bigquery.TableFieldSchema(
name='s', type='STRING', mode='NULLABLE', description='s description')
@@ -109,6 +110,7 @@ class TestTableSchemaParser(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestTableReferenceParser(unittest.TestCase):
+
def test_calling_with_table_reference(self):
table_ref = bigquery.TableReference()
table_ref.projectId = 'test_project'
@@ -175,6 +177,7 @@ class TestTableReferenceParser(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestBigQueryWrapper(unittest.TestCase):
+
def test_delete_non_existing_dataset(self):
client = mock.Mock()
client.datasets.Delete.side_effect = HttpError(
@@ -388,6 +391,7 @@ class TestBigQueryWrapper(unittest.TestCase):
False)
def test_wait_for_job_returns_true_when_job_is_done(self):
+
def make_response(state):
m = mock.Mock()
m.status.errorResult = None
@@ -581,6 +585,7 @@ class TestBigQueryWrapper(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestRowAsDictJsonCoder(unittest.TestCase):
+
def test_row_as_dict(self):
coder = RowAsDictJsonCoder()
test_value = {'s': 'abc', 'i': 123, 'f': 123.456, 'b': True}
@@ -622,6 +627,7 @@ class TestRowAsDictJsonCoder(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestJsonRowWriter(unittest.TestCase):
+
def test_write_row(self):
rows = [
{
@@ -654,6 +660,7 @@ class TestJsonRowWriter(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestAvroRowWriter(unittest.TestCase):
+
def test_write_row(self):
schema = bigquery.TableSchema(
fields=[
@@ -682,6 +689,7 @@ class TestAvroRowWriter(unittest.TestCase):
class TestBQJobNames(unittest.TestCase):
+
def test_simple_names(self):
self.assertEqual(
"beam_bq_job_EXPORT_beamappjobtest_abcd",
@@ -720,6 +728,7 @@ class TestBQJobNames(unittest.TestCase):
@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
class TestCheckSchemaEqual(unittest.TestCase):
+
def test_simple_schemas(self):
schema1 = bigquery.TableSchema(fields=[])
self.assertTrue(check_schema_equal(schema1, schema1))
@@ -985,9 +994,8 @@ class TestBeamTypehintFromSchema(unittest.TestCase):
schema = {"fields": self.get_schema_fields_with_mode("repeated")}
typehints = get_beam_typehints_from_tableschema(schema)
- expected_repeated_typehints = [
- (name, Sequence[type]) for name, type in self.EXPECTED_TYPEHINTS
- ]
+ expected_repeated_typehints = [(name, Sequence[type])
+ for name, type in self.EXPECTED_TYPEHINTS]
self.assertEqual(typehints, expected_repeated_typehints)
@@ -995,9 +1003,8 @@ class TestBeamTypehintFromSchema(unittest.TestCase):
schema = {"fields": self.get_schema_fields_with_mode("nullable")}
typehints = get_beam_typehints_from_tableschema(schema)
- expected_nullable_typehints = [
- (name, Optional[type]) for name, type in self.EXPECTED_TYPEHINTS
- ]
+ expected_nullable_typehints = [(name, Optional[type])
+ for name, type in self.EXPECTED_TYPEHINTS]
self.assertEqual(typehints, expected_nullable_typehints)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py b/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
index cd3edf19de5..78393844787 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_write_it_test.py
@@ -591,8 +591,10 @@ class BigQueryWriteIntegrationTests(unittest.TestCase):
max_file_size=1, # bytes
method=beam.io.WriteToBigQuery.Method.FILE_LOADS,
additional_bq_parameters={
- 'schemaUpdateOptions': ['ALLOW_FIELD_ADDITION',
- 'ALLOW_FIELD_RELAXATION']},
+ 'schemaUpdateOptions': [
+ 'ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION'
+ ]
+ },
temp_file_format=file_format))
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py b/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py
index 1aafb1b60a8..f19ba97048d 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_write_perf_test.py
@@ -67,6 +67,7 @@ from apache_beam.testing.synthetic_pipeline import SyntheticSource
class BigQueryWritePerfTest(LoadTest):
+
def __init__(self):
super().__init__()
self.output_dataset = self.pipeline.get_option('output_dataset')
diff --git a/sdks/python/apache_beam/io/gcp/bigtableio.py b/sdks/python/apache_beam/io/gcp/bigtableio.py
index ffb1852eb0f..1911f62e2fc 100644
--- a/sdks/python/apache_beam/io/gcp/bigtableio.py
+++ b/sdks/python/apache_beam/io/gcp/bigtableio.py
@@ -80,6 +80,7 @@ class _BigTableWriteFn(beam.DoFn):
table_id(str): GCP Table ID
"""
+
def __init__(self, project_id, instance_id, table_id):
""" Constructor of the Write connector of Bigtable
Args:
@@ -253,6 +254,7 @@ class WriteToBigTable(beam.PTransform):
self._project_id, self._instance_id, self._table_id)))
class _DirectRowMutationsToBeamRow(beam.DoFn):
+
def process(self, direct_row):
args = {"key": direct_row.row_key, "mutations": []}
# start accumulating mutations in a list
@@ -346,6 +348,7 @@ class ReadFromBigtable(PTransform):
# To make use of those methods and to give Python users a more familiar
# object, we process each Beam Row and return a PartialRowData equivalent.
class _BeamRowToPartialRowData(beam.DoFn):
+
def process(self, row):
key = row.key
families = row.column_families
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py
index f56443fe6ed..491213a86da 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastore_write_it_pipeline.py
@@ -66,6 +66,7 @@ class EntityWrapper(object):
Namespace and project are taken from the parent key.
"""
+
def __init__(self, kind, parent_key):
self._kind = kind
self._parent_key = parent_key
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
index f120234e974..bc609e52427 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio.py
@@ -173,6 +173,7 @@ class ReadFromDatastore(PTransform):
@typehints.with_output_types(types.Query)
class _SplitQueryFn(DoFn):
"""A `DoFn` that splits a given query into multiple sub-queries."""
+
def __init__(self, num_splits):
super().__init__()
self._num_splits = num_splits
@@ -281,6 +282,7 @@ class ReadFromDatastore(PTransform):
@typehints.with_output_types(types.Entity)
class _QueryFn(DoFn):
"""A DoFn that fetches entities from Cloud Datastore, for a given query."""
+
def process(self, query, *unused_args, **unused_kwargs):
if query.namespace is None:
query.namespace = ''
@@ -359,6 +361,7 @@ class _Mutate(PTransform):
should be idempotent (`upsert` and `delete` mutations) to prevent duplicate
data or errors.
"""
+
def __init__(self, project):
"""
Args:
@@ -523,6 +526,7 @@ class WriteToDatastore(_Mutate):
property key is empty then it is filled with the project ID passed to this
transform.
"""
+
def __init__(
self,
project,
@@ -540,6 +544,7 @@ class WriteToDatastore(_Mutate):
super().__init__(mutate_fn, throttle_rampup, hint_num_workers)
class _DatastoreWriteFn(_Mutate.DatastoreMutateFn):
+
def element_to_client_batch_item(self, element):
if not isinstance(element, types.Entity):
raise ValueError(
@@ -575,6 +580,7 @@ class DeleteFromDatastore(_Mutate):
project ID passed to this transform. If ``project`` field in key is empty then
it is filled with the project ID passed to this transform.
"""
+
def __init__(
self,
project,
@@ -593,6 +599,7 @@ class DeleteFromDatastore(_Mutate):
super().__init__(mutate_fn, throttle_rampup, hint_num_workers)
class _DatastoreDeleteFn(_Mutate.DatastoreMutateFn):
+
def element_to_client_batch_item(self, element):
if not isinstance(element, types.Key):
raise ValueError(
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
index aac99cb8c1f..e15895c1d02 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/datastoreio_test.py
@@ -51,6 +51,7 @@ except ImportError:
# used for internal testing only
class FakeMessage:
+
def __init__(self, entity, key):
self.entity = entity
self.key = key
@@ -64,6 +65,7 @@ class FakeMessage:
# used for internal testing only
class FakeMutation(object):
+
def __init__(self, entity=None, key=None):
"""Fake mutation request object.
@@ -81,6 +83,7 @@ class FakeMutation(object):
class FakeBatch(object):
+
def __init__(self, all_batch_items=None, commit_count=None):
"""Fake ``google.cloud.datastore.batch.Batch`` object.
@@ -116,6 +119,7 @@ class FakeBatch(object):
@unittest.skipIf(client is None, 'Datastore dependencies are not installed')
class MutateTest(unittest.TestCase):
+
def test_write_mutations_no_errors(self):
mock_batch = MagicMock()
mock_throttler = MagicMock()
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py
index 842579dfb40..a9c3be668bd 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/query_splitter.py
@@ -133,6 +133,7 @@ class IdOrName(object):
Implements sort ordering: by ID, then by name, keys with IDs before those
with names.
"""
+
def __init__(self, id_or_name):
self.id_or_name = id_or_name
if isinstance(id_or_name, str):
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn_test.py
index 0bbe953e93d..a0baab36542 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn_test.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/rampup_throttling_fn_test.py
@@ -33,6 +33,7 @@ class _RampupDelayException(Exception):
class RampupThrottlerTransformTest(unittest.TestCase):
+
@patch('datetime.datetime')
@patch('time.sleep')
def test_rampup_throttling(self, mock_sleep, mock_datetime):
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
index f7ce69099ca..98a7fca27f6 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/types.py
@@ -37,6 +37,7 @@ __all__ = ['Query', 'Key', 'Entity']
class Query(object):
+
def __init__(
self,
kind=None,
@@ -152,6 +153,7 @@ class Query(object):
class Key(object):
+
def __init__(
self,
path_elements: List[Union[str, int]],
@@ -228,6 +230,7 @@ class Key(object):
class Entity(object):
+
def __init__(self, key: Key, exclude_from_indexes: Iterable[str] = ()):
"""
Represents a Datastore entity.
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/util.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/util.py
index 06a22143f59..213f289f26b 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/util.py
+++ b/sdks/python/apache_beam/io/gcp/datastore/v1new/util.py
@@ -46,6 +46,7 @@ class MovingSum(object):
convenience we expose the count of entries as well so this doubles as a
moving average tracker.
"""
+
def __init__(self, window_ms, bucket_ms):
if window_ms < bucket_ms or bucket_ms <= 0:
raise ValueError("window_ms >= bucket_ms > 0 please")
@@ -110,6 +111,7 @@ class MovingSum(object):
class DynamicBatchSizer(object):
"""Determines request sizes for future Datastore RPCs."""
+
def __init__(self):
self._commit_time_per_entity_ms = MovingSum(
window_ms=120000, bucket_ms=10000)
diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1new/util_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1new/util_test.py
index f82b223b67e..de97c1d699e 100644
--- a/sdks/python/apache_beam/io/gcp/datastore/v1new/util_test.py
... 27217 lines suppressed ...