Merge pull request #121 from V0RT3X4/allow-all-cols

fix: CargoMovementsResult to_df all columns, closes #120
VorTECHsa · Dec 5, 2019 · 01115f2 · 01115f2
2 parents 80897e7 + 98af1ef
commit 01115f2
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 1 deletion.
diff --git a/tests/endpoints/test_cargo_movements_real.py b/tests/endpoints/test_cargo_movements_real.py
@@ -46,6 +46,21 @@ def test_search_single_filter_id(self):
 
         assert len(df) == 2
 
+    def test_to_df_all_columns(self):
+        df = (
+            CargoMovements()
+            .search(
+                filter_activity="loading_state",
+                filter_products="6f11b0724c9a4e85ffa7f1445bc768f054af755a090118dcf99f14745c261653",
+                filter_time_min=datetime(2019, 8, 29),
+                filter_time_max=datetime(2019, 8, 29, 0, 10),
+            )
+            .to_df(columns="all")
+            .head(2)
+        )
+
+        assert len(df) == 2
+
     def test_search_single_filter_origin_name(self):
         df = (
             CargoMovements()

diff --git a/tests/endpoints/test_vessel_movements_real.py b/tests/endpoints/test_vessel_movements_real.py
@@ -15,6 +15,20 @@ def test_search(self):
 
         assert len(v) > 10
 
+    def test_to_df_all_columns(self):
+        df = (
+            VesselMovements()
+            .search(
+                filter_time_min=datetime(2017, 10, 1, 0, 0),
+                filter_time_max=datetime(2017, 10, 1, 0, 10),
+                filter_origins="rotterdam",
+            )
+            .to_df(columns="all")
+            .head(2)
+        )
+
+        assert len(df) == 2
+
     def test_search_to_dataframe(self):
         df = (
             VesselMovements()

diff --git a/vortexasdk/endpoints/cargo_movements_result.py b/vortexasdk/endpoints/cargo_movements_result.py
@@ -543,7 +543,10 @@ def to_df(self, columns=None) -> pd.DataFrame:
         with Pool(os.cpu_count()) as pool:
             records = pool.map(flatten, super().to_list())
 
-        return pd.DataFrame(data=records, columns=columns)
+        if columns == "all":
+            return pd.DataFrame(data=records)
+        else:
+            return pd.DataFrame(data=records, columns=columns)
 
 
 DEFAULT_COLUMNS = [