From 75739536b32ae7bd67f6f23fcf6bc8ee76e833da Mon Sep 17 00:00:00 2001
From: Kael Zhang <i+github@kael.me>
Date: Thu, 12 Dec 2024 18:01:58 +0800
Subject: [PATCH] fixes #23: supports the Donchian Channels

---
 docs/README.md                               | 29 +++++++++++
 stock_pandas/commands/over_bought_or_sold.py | 55 ++++++++++++++++----
 test/test_commands.py                        |  8 +++
 3 files changed, 81 insertions(+), 11 deletions(-)

diff --git a/docs/README.md b/docs/README.md
index 3bda410..800ea7c 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -655,6 +655,35 @@ hhv:<period>,<column>
 
 Gets the highest of high prices in N periods. The arguments of `hhv` is the same as `llv`
 
+### `donchian`, Donchian Channels
+
+```
+donchian:<period>,<column_upper>,<column_lower>
+donchian.upper:<period>,<column_upper>
+donchian.lower:<period>,<column_lower>
+```
+
+Gets the Donchian channels
+
+- **period** `int`
+- **column_upper?** `str='high'` The column to calculate highest high values, defaults to `'high'`
+- **column_lower?** `str='low'` The column to calculate lowest low values, defaults to `'low'`
+
+```py
+# Donchian middle channel
+stock['donchian']
+stock['donchian.middle']
+
+# Donchian upper channel, a shortcut for stock['donchian.upper']
+stock['donchian.u']
+stock['donchian.upper']
+
+# Donchian lower channel, which is equivalent to stock['donchian.l']
+stock['donchian.lower']
+stock['donchian.l']
+```
+
+
 ## Built-in Commands for Statistics
 
 ### `column`
diff --git a/stock_pandas/commands/over_bought_or_sold.py b/stock_pandas/commands/over_bought_or_sold.py
index 9d709b9..4b68ae0 100644
--- a/stock_pandas/commands/over_bought_or_sold.py
+++ b/stock_pandas/commands/over_bought_or_sold.py
@@ -36,15 +36,20 @@ def llv(df, s, period, column) -> ReturnType:
         min
     ), period
 
+arg_column_low = ('low', column_enums)
+llv_args = [
+    arg_period,
+    arg_column_low
+]
+
+arg_column_high = ('high', column_enums)
+hhv_args = [
+    arg_period,
+    arg_column_high
+]
 
 COMMANDS['llv'] = (
-    CommandPreset(
-        llv,
-        [
-            arg_period,
-            ('low', column_enums)
-        ]
-    ),
+    CommandPreset(llv, llv_args),
     None,
     None
 )
@@ -62,15 +67,43 @@ def hhv(df, s, period, column) -> ReturnType:
 
 
 COMMANDS['hhv'] = (
+    CommandPreset(hhv, hhv_args),
+    None,
+    None
+)
+
+
+# Donchian Channel
+# ref: https://en.wikipedia.org/wiki/Donchian_channel
+
+def donchian(df, s, period, hhv_column, llv_column) -> ReturnType:
+    """Gets Donchian Channel
+    """
+
+    hhv = df.exec(f'hhv:{period},{hhv_column}')[s]
+    llv = df.exec(f'llv:{period},{llv_column}')[s]
+
+    return (hhv + llv) / 2, period
+
+
+COMMANDS['donchian'] = (  # type: ignore
     CommandPreset(
-        hhv,
+        donchian,
         [
             arg_period,
-            ('high', column_enums)
+            arg_column_high,
+            arg_column_low
         ]
     ),
-    None,
-    None
+    dict(
+        upper=CommandPreset(hhv, hhv_args),
+        lower=CommandPreset(llv, llv_args)
+    ),
+    dict(
+        u='upper',
+        l='lower',
+        middle=None
+    )
 )
 
 
diff --git a/test/test_commands.py b/test/test_commands.py
index 6a5f19a..015e099 100644
--- a/test/test_commands.py
+++ b/test/test_commands.py
@@ -89,6 +89,14 @@ def test_boll(stock):
     assert stock['column:close > boll']['2020-02-05']
 
 
+def test_donchian(stock):
+    assert stock['donchian.upper:20'].equals(stock['hhv:20'])
+    assert stock['donchian.lower:20'].equals(stock['llv:20'])
+    assert stock['donchian:20'].equals(
+        (stock['hhv:20'] + stock['llv:20']) / 2
+    )
+
+
 def test_macd(stock):
     assert stock['macd / macd.dea']['2020-02-10']
     assert stock['macd.histogram / 0']['2020-02-10']