Adjust type coercion (#187)

earthobservations · Sep 25, 2020 · 5b6ea70 · 5b6ea70
1 parent 0e001cf
commit 5b6ea70
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 13 deletions.
diff --git a/tests/additionals/test_functions.py b/tests/additionals/test_functions.py
@@ -3,6 +3,7 @@
 import pytest
 import numpy as np
 import pandas as pd
+from pandas._testing import assert_frame_equal
 
 from wetterdienst.additionals.functions import (
     check_parameters,
@@ -50,18 +51,17 @@ def test_coerce_field_types():
 
     expected_df = pd.DataFrame(
         {
-            "QN": pd.Series([1], dtype=np.int32),
-            "RS_IND_01": pd.Series([1], dtype=np.int32),
+            "QN": pd.Series([1], dtype=np.int8),
+            "RS_IND_01": pd.Series([1], dtype=np.int8),
             "DATE": [pd.Timestamp("1970-01-01")],
             "END_OF_INTERVAL": [pd.Timestamp("1970-01-01")],
             "V_VV_I": ["P"],
         }
     )
 
-    assert (
-        coerce_field_types(df, TimeResolution.HOURLY).values.tolist()
-        == expected_df.values.tolist()
-    )
+    df = coerce_field_types(df, TimeResolution.HOURLY)
+
+    assert_frame_equal(df, expected_df)
 
 
 def test_create_humanized_column_names_mapping():

diff --git a/wetterdienst/additionals/functions.py b/wetterdienst/additionals/functions.py
@@ -261,7 +261,6 @@ def coerce_field_types(
     """
 
     for column in df.columns:
-        column_value_index = df[column].notna()
 
         # Station ids are handled separately as they are expected to not have any nans
         if column == DWDMetaColumns.STATION_ID.value:
@@ -276,13 +275,9 @@ def coerce_field_types(
                 df[column], format=DatetimeFormat.YMDH_COLUMN_M.value
             )
         elif column in QUALITY_FIELDS or column in INTEGER_FIELDS:
-            df.loc[column_value_index, column] = df.loc[
-                column_value_index, column
-            ].astype(int)
+            df[column] = df[column].astype("int8", errors="ignore")
         elif column in STRING_FIELDS:
-            df.loc[column_value_index, column] = df.loc[
-                column_value_index, column
-            ].astype(str)
+            df[column] = df[column].astype("str")
         else:
             df[column] = df[column].astype(float)