refactor: apply refurb suggestions

kraina-ai · Jan 2, 2024 · 1ffdf78 · 1ffdf78
1 parent 67a686e
commit 1ffdf78
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 14 deletions.
diff --git a/quackosm/functions.py b/quackosm/functions.py
@@ -212,13 +212,12 @@ def convert_pbf_to_gpq(
         │ 2140 rows (20 shown)                                                         3 columns │
         └────────────────────────────────────────────────────────────────────────────────────────┘
     """
-    reader = PbfFileReader(
+    return PbfFileReader(
         tags_filter=tags_filter,
         geometry_filter=geometry_filter,
         working_directory=working_directory,
         osm_way_polygon_features_config=osm_way_polygon_features_config,
-    )
-    return reader.convert_pbf_to_gpq(
+    ).convert_pbf_to_gpq(
         pbf_path=pbf_path,
         result_file_path=result_file_path,
         explode_tags=explode_tags,
@@ -386,13 +385,12 @@ def get_features_gdf(
         <BLANKLINE>
         [3109 rows x 2 columns]
     """
-    reader = PbfFileReader(
+    return PbfFileReader(
         tags_filter=tags_filter,
         geometry_filter=geometry_filter,
         working_directory=working_directory,
         osm_way_polygon_features_config=osm_way_polygon_features_config,
-    )
-    return reader.get_features_gdf(
+    ).get_features_gdf(
         file_paths=file_paths,
         explode_tags=explode_tags,
         ignore_cache=ignore_cache,

diff --git a/quackosm/pbf_file_reader.py b/quackosm/pbf_file_reader.py
@@ -1156,7 +1156,7 @@ def _save_parquet_file_with_geometry(
         fix_geometries: bool = False,
     ) -> "duckdb.DuckDBPyRelation":
         if not fix_geometries:
-            with TaskProgressSpinner(step_name, f"{step_number}"):
+            with TaskProgressSpinner(step_name, str(step_number)):
                 self.connection.sql(f"""
                     COPY (
                         SELECT
@@ -1306,8 +1306,7 @@ def _concatenate_results_to_geoparquet(
 
         valid_features_parquet_table = valid_features_parquet_table.append_column(
             GEOMETRY_COLUMN, geometry_column
-        )
-        valid_features_parquet_table = valid_features_parquet_table.drop("geometry_wkb")
+        ).drop("geometry_wkb")
 
         parquet_tables = [valid_features_parquet_table]
 

diff --git a/tests/test_pbf_file_reader.py b/tests/test_pbf_file_reader.py
@@ -81,8 +81,7 @@ def test_pbf_reader(
     expected_features_columns_length: int,
 ):
     """Test proper files loading in `PbfFileReader`."""
-    handler = PbfFileReader(tags_filter=query)
-    features_gdf = handler.get_features_gdf(
+    features_gdf = PbfFileReader(tags_filter=query).get_features_gdf(
         file_paths=[Path(__file__).parent / "test_files" / test_file_name],
         explode_tags=True,
         ignore_cache=True,
@@ -99,10 +98,9 @@ def test_pbf_reader(
 def test_pbf_reader_geometry_filtering():  # type: ignore
     """Test proper spatial data filtering in `PbfFileReader`."""
     file_name = "d17f922ed15e9609013a6b895e1e7af2d49158f03586f2c675d17b760af3452e.osm.pbf"
-    handler = PbfFileReader(
+    features_gdf = PbfFileReader(
         tags_filter=HEX2VEC_FILTER, geometry_filter=Polygon([(0, 0), (0, 1), (1, 1), (1, 0)])
-    )
-    features_gdf = handler.get_features_gdf(
+    ).get_features_gdf(
         file_paths=[Path(__file__).parent / "test_files" / file_name],
         explode_tags=True,
         ignore_cache=True,