refactor(response): simplify time parsing (#160)

Co-authored-by: Matthias (~talfus-laddus) <[email protected]>
GIScience · Jul 8, 2024 · 435066d · 435066d
1 parent ebcd7ee
commit 435066d
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 50 deletions.
diff --git a/ohsome/response.py b/ohsome/response.py
@@ -67,7 +67,11 @@ def _as_dataframe(self, multi_index=True):
         else:
             raise TypeError("This result type is not implemented.")
 
-        self._format_timestamp(result_df)
+        time_columns = result_df.columns.intersection(
+            ["timestamp", "fromTimestamp", "toTimestamp"]
+        )
+        result_df[time_columns] = result_df[time_columns].apply(self._format_timestamp)
+
         if multi_index:
             self._set_index(result_df, groupby_names)
 
@@ -110,37 +114,17 @@ def _as_geodataframe(
                 "This result type cannot be converted to a GeoPandas GeoDataFrame object."
             )
 
-        if "@validFrom" in features.columns:
-            features["@validFrom"] = pd.to_datetime(
-                features["@validFrom"].str.replace("Z", ""), format="ISO8601"
-            )
-            features["@validTo"] = pd.to_datetime(
-                features["@validTo"].str.replace("Z", ""), format="ISO8601"
-            )
-            if multi_index:
-                features = features.set_index(["@osmId", "@validFrom", "@validTo"])
-        elif "@snapshotTimestamp" in features.columns:
-            features["@snapshotTimestamp"] = pd.to_datetime(
-                features["@snapshotTimestamp"].str.replace("Z", ""), format="ISO8601"
-            )
-            if multi_index:
-                features = features.set_index(["@osmId", "@snapshotTimestamp"])
-        elif (
-            "timestamp" in features.columns and "groupByBoundaryId" in features.columns
-        ):
-            features["timestamp"] = pd.to_datetime(
-                features["timestamp"].str.replace("Z", ""), format="ISO8601"
-            )
-            if multi_index:
-                features = features.set_index(["groupByBoundaryId", "timestamp"])
-        elif "@timestamp" in features.columns:
-            features["@timestamp"] = pd.to_datetime(
-                features["@timestamp"].str.replace("Z", ""), format="ISO8601"
-            )
-            if multi_index:
-                features = features.set_index(["@timestamp"])
-        else:
-            raise TypeError("This result type is not implemented.")
+        time_columns = ["@validFrom", "@validTo", "@snapshotTimestamp", "@timestamp"]
+        existing_time_columns = features.columns.intersection(time_columns)
+        features[existing_time_columns] = features[existing_time_columns].apply(
+            self._format_timestamp
+        )
+
+        if multi_index:
+            index_columns = features.columns.intersection(
+                ["@osmId"] + time_columns
+            ).to_list()
+            features = features.set_index(index_columns)
 
         return features.sort_index()
 
@@ -192,20 +176,7 @@ def _create_groupby_dataframe(self, data, groupby_names) -> DataFrame:
                 record_dfs.extend(record_result)
         return pd.DataFrame().from_records(record_dfs)
 
-    def _format_timestamp(self, result_df: DataFrame) -> None:
-        """
-        Format timestamp column as datetime
-        :param result_df:
-        :return:
-        """
-        if "timestamp" in result_df.columns:
-            result_df["timestamp"] = pd.to_datetime(
-                result_df["timestamp"].str.replace("Z", ""), format="ISO8601"
-            )
-        else:
-            result_df["fromTimestamp"] = pd.to_datetime(
-                result_df["fromTimestamp"].str.replace("Z", ""), format="ISO8601"
-            )
-            result_df["toTimestamp"] = pd.to_datetime(
-                result_df["toTimestamp"].str.replace("Z", ""), format="ISO8601"
-            )
+    @staticmethod
+    def _format_timestamp(dt: pd.Series) -> pd.Series:
+        """Format timestamp column as datetime."""
+        return pd.to_datetime(dt.str.replace("Z", ""), format="ISO8601")
diff --git a/ohsome/test/test_response.py b/ohsome/test/test_response.py
@@ -443,7 +443,7 @@ def test_all_columns_with_timestamps_to_be_without_timezone(base_client):
     at_timestamp = (
         client.contributions.geometry.post(time=time2iso, bboxes=bbox, filter=fltr)
         .as_dataframe()
-        .index[0]
+        .index[0][1]
     )
     timestamp = (
         client.elements.count.groupByBoundary.post(bboxes=bbox, time=time, filter=fltr)