improved parse_load to handle different process types correctly

fixed #125
EnergieID · Aug 28, 2021 · 9fc9b0c · 9fc9b0c
1 parent a33dc1b
commit 9fc9b0c
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 13 deletions.
diff --git a/entsoe/entsoe.py b/entsoe/entsoe.py
@@ -1067,7 +1067,7 @@ def query_load(self, country_code: Union[Area, str], start: pd.Timestamp,
     @year_limited
     def query_load_forecast(
             self, country_code: Union[Area, str], start: pd.Timestamp,
-            end: pd.Timestamp, process_type: str = 'A01') -> pd.Series:
+            end: pd.Timestamp, process_type: str = 'A01') -> pd.DataFrame:
         """
         Parameters
         ----------
@@ -1078,15 +1078,17 @@ def query_load_forecast(
 
         Returns
         -------
-        pd.Series
+        pd.DataFrame
         """
         area = lookup_area(country_code)
         text = super(EntsoePandasClient, self).query_load_forecast(
             country_code=area, start=start, end=end, process_type=process_type)
-        series = parse_loads(text)
-        series = series.tz_convert(area.tz)
-        series = series.truncate(before=start, after=end)
-        return series
+
+        df = parse_loads(text, process_type=process_type)
+        df = df.tz_convert(area.tz)
+        df = df.truncate(before=start, after=end)
+        return df
+
 
     @year_limited
     def query_generation_forecast(

diff --git a/entsoe/parsers.py b/entsoe/parsers.py
@@ -63,21 +63,40 @@ def parse_netpositions(xml_text):
     return series
 
 
-def parse_loads(xml_text):
+def parse_loads(xml_text, process_type='A01'):
     """
     Parameters
     ----------
     xml_text : str
 
     Returns
     -------
-    pd.Series
+    pd.DataFrame
     """
-    series = pd.Series(dtype = 'object')
-    for soup in _extract_timeseries(xml_text):
-        series = series.append(_parse_load_timeseries(soup))
-    series = series.sort_index()
-    return series
+    if process_type == 'A01' or process_type == 'A16':
+        series = pd.Series(dtype = 'object')
+        for soup in _extract_timeseries(xml_text):
+                series = series.append(_parse_load_timeseries(soup))
+        series = series.sort_index()
+        return pd.DataFrame({
+            'Forecasted Load' if process_type == 'A01' else 'Actual Load': series
+        })
+    else:
+        series_min = pd.Series(dtype='object')
+        series_max = pd.Series(dtype='object')
+        for soup in _extract_timeseries(xml_text):
+            t = _parse_load_timeseries(soup)
+            if soup.find('businesstype').text == 'A60':
+                series_min = series_min.append(t)
+            elif soup.find('businesstype').text == 'A61':
+                series_max = series_max.append(t)
+            else:
+                continue
+        return pd.DataFrame({
+            'Min Forecasted Load': series_min,
+            'Max Forecasted Load': series_max
+        })
+
 
 
 def parse_generation(