Update location code argument validation

geomagio/Controller.py

@@ -566,9 +566,11 @@ def parse_args(args):
             metavar=('FROM', 'TO'),
             nargs=2)
     parser.add_argument('--locationcode',
-            choices=['R0', 'R1', 'RM', 'Q0', 'D0', 'C0'])
+            help='EDGE location code, e.g. "R0", "R1"',
+            type=edge.LocationCode)
     parser.add_argument('--outlocationcode',
-            choices=['R0', 'R1', 'RM', 'Q0', 'D0', 'C0'])
+            help='EDGE output location code (when different from --locationcode)',
+            type=edge.LocationCode)
     parser.add_argument('--interval',
             default='minute',
             choices=['hourly', 'minute', 'second'])

geomagio/edge/LocationCode.py

+"""EDGE Location Code argument validation."""
+
+import argparse
+import re
+
+
+def LocationCode(code):
+    """EDGE Location Code argument validator.
+
+    Location Code is the last component in a channel identifier;
+    SNCL => Station, Network, Channel, Location Code
+
+    Parameters
+    ----------
+    code : str
+        the code to validate
+
+    Returns
+    -------
+    str
+        validated lcoation code.
+
+    Raises
+    ------
+    argparse.ArgumentTypeError
+        if the location code doesn't match the regular expression.
+    """
+    try:
+        return re.match('^[A-Z0-9]{2}$', code).group(0)
+    except:
+        raise argparse.ArgumentTypeError('Invalid location code, expected /^[A-Z0-9]{2}$/')

geomagio/edge/__init__.py

@@ -2,9 +2,11 @@
 """
 
 from EdgeFactory import EdgeFactory
+from LocationCode import LocationCode
 from RawInputClient import RawInputClient
 
 __all__ = [
     'EdgeFactory',
+    'LocationCode',
     'RawInputClient'
 ]