From 6fe61647028e0060691aded5fef0be464c548aa5 Mon Sep 17 00:00:00 2001
From: Nicolas Schmid <nicolas@breiten.ch>
Date: Tue, 15 Aug 2023 17:25:11 +0200
Subject: [PATCH] feat: only add additional magnitudes if includeallmagnitudes
 is requested

---
 catalog_tools/io/client.py | 16 ++++++++++------
 catalog_tools/io/parser.py | 16 ++++++++++------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/catalog_tools/io/client.py b/catalog_tools/io/client.py
index a2a7fae..23ef3ec 100644
--- a/catalog_tools/io/client.py
+++ b/catalog_tools/io/client.py
@@ -6,7 +6,7 @@
 import requests
 
 from catalog_tools.catalog import Catalog
-from catalog_tools.io.parser import CustomContentHandler
+from catalog_tools.io.parser import QuakeMLHandler
 
 
 class FDSNWSEventClient():
@@ -28,7 +28,8 @@ def get_events(self, start_time: Optional[datetime] = None,
                    min_magnitude: Optional[float] = None,
                    include_all_magnitudes: Optional[bool] = None,
                    event_type: Optional[str] = None,
-                   delta_m: float = 0.1) -> pd.DataFrame:
+                   delta_m: float = 0.1,
+                   include_uncertainty: bool = False) -> pd.DataFrame:
         """Downloads an earthquake catalog based on a URL.
 
         Args:
@@ -73,17 +74,20 @@ def get_events(self, start_time: Optional[datetime] = None,
         if event_type:
             request_url += f'&eventtype={event_type}'
 
-        print(request_url)
-
         catalog = []
         parser = make_parser()
         parser.setFeature(handler.feature_namespaces, False)
-        parser.setContentHandler(CustomContentHandler(catalog))
+        parser.setContentHandler(QuakeMLHandler(
+            catalog, includeallmagnitudes=include_all_magnitudes))
 
         r = requests.get(request_url, stream=True)
         r.raw.decode_content = True  # if content-encoding is used decode
         parser.parse(r.raw)
 
         df = Catalog.from_dict(catalog)
-
+        # if not include_uncertainty:
+        # wrong way round
+        # rgx = "(_uncertainty|_lowerUncertainty|" \
+        #       "_upperUncertainty|_confidenceLevel)$"
+        # df = df.filter(regex=rgx)
         return df
diff --git a/catalog_tools/io/parser.py b/catalog_tools/io/parser.py
index 3e01311..9b82fec 100644
--- a/catalog_tools/io/parser.py
+++ b/catalog_tools/io/parser.py
@@ -117,7 +117,8 @@ def extract_secondary_magnitudes(magnitudes: list) -> dict:
     return magnitude_dict
 
 
-def parse_to_dict(event: dict, origins: list, magnitudes: list) -> dict:
+def parse_to_dict(event: dict, origins: list, magnitudes: list,
+                  includeallmagnitudes: bool = True) -> dict:
     preferred_origin = \
         get_preferred_origin(origins,
                              event.get('preferredOriginID', None))
@@ -126,8 +127,10 @@ def parse_to_dict(event: dict, origins: list, magnitudes: list) -> dict:
         get_preferred_magnitude(magnitudes,
                                 event.get('preferredMagnitudeID', None))
 
-    if magnitudes:
+    if magnitudes and includeallmagnitudes:
         magnitudes = select_secondary_magnitudes(magnitudes)
+    else:
+        magnitudes = []
 
     event_params = \
         {value: event.get(key, None) for key, value in EVENT_MAPPINGS.items()}
@@ -138,15 +141,15 @@ def parse_to_dict(event: dict, origins: list, magnitudes: list) -> dict:
         extract_secondary_magnitudes(magnitudes)
 
 
-class CustomContentHandler(xml.sax.ContentHandler):
+class QuakeMLHandler(xml.sax.ContentHandler):
     """
     Custom ContentHandler class that extends ContenHandler to
     stream parse QuakeML files.
     """
 
-    def __init__(self, catalog):
+    def __init__(self, catalog, includeallmagnitudes=True):
         self.catalog = catalog
-
+        self.includeallmagnitudes = includeallmagnitudes
         self.event = []
         self.origin = []
         self.magnitude = []
@@ -176,7 +179,8 @@ def startElement(self, tagName, attrs):
     def endElement(self, tagName):
         if tagName == 'event':
             self.catalog.append(parse_to_dict(
-                self.event[-1], self.origin, self.magnitude))
+                self.event[-1], self.origin, self.magnitude,
+                includeallmagnitudes=self.includeallmagnitudes))
             self.parent = ''
             self.location = ''
             self.event = []