feat: accept string aliases for method/extmethod constants (closes #14)

CHANGELOG.md

@@ -3,6 +3,10 @@ Change Log
 
 0.1.12
 -------
+- `calc_support_resistance` and `get_extrema` now accept string names for their
+  integer method constants (e.g. ``method='METHOD_NCUBED'`` in addition to
+  ``method=METHOD_NCUBED``); invalid strings raise a clear ``ValueError``
+  (closes #14)
 - Added `pandas_to_ohlc(df, low_col=None, high_col=None, close_col=None)` helper
   that converts any OHLC pandas DataFrame (yfinance, ccxt, or custom) into the
   ``(low_series, high_series)`` tuple expected by trendln functions;

tests/test_trendln.py

@@ -281,6 +281,78 @@ def test_tuple_both_none_raises(self):
 
 
 # ---------------------------------------------------------------------------
+# String aliases for method / extmethod constants (issue #14)
+# ---------------------------------------------------------------------------
+
+class TestStringMethodAliases:
+    """calc_support_resistance and get_extrema accept string names for their
+    integer method constants, e.g. 'METHOD_NCUBED' as well as METHOD_NCUBED."""
+
+    # --- method parameter ---
+
+    def test_method_string_ncubed(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, method='METHOD_NCUBED'),
+            calc_support_resistance(DATA_SIMPLE, method=METHOD_NCUBED),
+        )
+
+    def test_method_string_nsquredlogn(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, method='METHOD_NSQUREDLOGN'),
+            calc_support_resistance(DATA_SIMPLE, method=METHOD_NSQUREDLOGN),
+        )
+
+    def test_method_string_houghpoints(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, method='METHOD_HOUGHPOINTS'),
+            calc_support_resistance(DATA_SIMPLE, method=METHOD_HOUGHPOINTS),
+        )
+
+    def test_method_string_houghlines(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, method='METHOD_HOUGHLINES'),
+            calc_support_resistance(DATA_SIMPLE, method=METHOD_HOUGHLINES),
+        )
+
+    def test_method_string_probhough(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, method='METHOD_PROBHOUGH'),
+            calc_support_resistance(DATA_SIMPLE, method=METHOD_PROBHOUGH),
+        )
+
+    # --- extmethod parameter ---
+
+    def test_extmethod_string_naive(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, extmethod='METHOD_NAIVE'),
+            calc_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE),
+        )
+
+    def test_extmethod_string_naiveconsec(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, extmethod='METHOD_NAIVECONSEC'),
+            calc_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVECONSEC),
+        )
+
+    def test_extmethod_string_numdiff(self):
+        assert_result(
+            calc_support_resistance(DATA_SIMPLE, extmethod='METHOD_NUMDIFF'),
+            calc_support_resistance(DATA_SIMPLE, extmethod=METHOD_NUMDIFF),
+        )
+
+    def test_get_extrema_string_extmethod(self):
+        assert get_extrema(DATA_SIMPLE, extmethod='METHOD_NAIVE') == \
+               get_extrema(DATA_SIMPLE, extmethod=METHOD_NAIVE)
+
+    # --- invalid string raises clearly ---
+
+    def test_invalid_method_string_raises(self):
+        with pytest.raises(ValueError, match='METHOD_BOGUS'):
+            calc_support_resistance(DATA_SIMPLE, method='METHOD_BOGUS')
+
+    def test_invalid_extmethod_string_raises(self):
+        with pytest.raises(ValueError, match='METHOD_BOGUS'):
+            get_extrema(DATA_SIMPLE, extmethod='METHOD_BOGUS')
 # pandas_to_ohlc
 # ---------------------------------------------------------------------------
 

trendln/__init__.py

@@ -395,6 +395,29 @@ def test_sup_res(curdir):
     return None
 METHOD_NAIVE, METHOD_NAIVECONSEC, METHOD_NUMDIFF = 0, 1, 2
 METHOD_NCUBED, METHOD_NSQUREDLOGN, METHOD_HOUGHPOINTS, METHOD_HOUGHLINES, METHOD_PROBHOUGH = 0, 1, 2, 3, 4
+_EXTMETHOD_NAMES = {
+    'METHOD_NAIVE': METHOD_NAIVE,
+    'METHOD_NAIVECONSEC': METHOD_NAIVECONSEC,
+    'METHOD_NUMDIFF': METHOD_NUMDIFF,
+}
+_METHOD_NAMES = {
+    'METHOD_NCUBED': METHOD_NCUBED,
+    'METHOD_NSQUREDLOGN': METHOD_NSQUREDLOGN,
+    'METHOD_HOUGHPOINTS': METHOD_HOUGHPOINTS,
+    'METHOD_HOUGHLINES': METHOD_HOUGHLINES,
+    'METHOD_PROBHOUGH': METHOD_PROBHOUGH,
+}
+def _resolve_name(value, lookup, param):
+    """Allow string aliases for integer method constants.
+    e.g. 'METHOD_NCUBED' is accepted in addition to METHOD_NCUBED.
+    """
+    if isinstance(value, str):
+        if value not in lookup:
+            raise ValueError(
+                f'{param} string {value!r} is not a valid name; '
+                f'use one of {list(lookup)} or the corresponding integer constant')
+        return lookup[value]
+    return value
 def check_num_alike(h):
     if type(h) is list and all([isinstance(x, (bool, int, float)) for x in h]): return True
     elif type(h) is np.ndarray and h.ndim==1 and h.dtype.kind in 'biuf': return True
@@ -492,6 +515,7 @@ def _resolve(explicit, default_key):
         'has columns named Low, High, or Close (case-insensitive).')
 
 def get_extrema(h, extmethod=METHOD_NUMDIFF, accuracy=2):
+    extmethod = _resolve_name(extmethod, _EXTMETHOD_NAMES, 'extmethod')
     #h must be single dimensional array-like object e.g. List, np.ndarray, pd.Series
     if type(h) is tuple and len(h) == 2 and (h[0] is None or check_num_alike(h[0])) and (h[1] is None or check_num_alike(h[1])) and (not h[0] is None or not h[1] is None):
         hmin, hmax = h[0], h[1]
@@ -556,6 +580,7 @@ def numdiff_extrema(func):
 def calc_support_resistance(h, extmethod = METHOD_NUMDIFF, method=METHOD_NSQUREDLOGN,
                             window=125, errpct=0.005, hough_scale=0.01, hough_prob_iter=10,
                             sortError=False, accuracy=2):
+    method = _resolve_name(method, _METHOD_NAMES, 'method')
     if not type(window) is int:
         raise ValueError('window must be of type int')
     if not type(errpct) is float: