feat: add ohlc candlestick rendering to plot functions (closes #29)

CHANGELOG.md

@@ -1,5 +1,15 @@
 Change Log
 ===========
+
+0.1.18
+-------
+- Added ``ohlc`` parameter (default ``None``) to ``plot_support_resistance`` and
+  ``plot_sup_res_date``; accepts a ``(opens, highs, lows, closes)`` 4-tuple of
+  numeric sequences and renders the price series as OHLC candlestick bars using
+  matplotlib Rectangle patches (green = bullish, red = bearish) while all
+  trendline overlays continue to work unchanged; ``hist`` still controls extrema
+  detection so the existing API is fully preserved (closes #29)
+
 0.1.17
 -------
 - Added ``get_horizontal_levels(h, pctbound=0.05, extmethod=..., accuracy=2,

setup.py

@@ -30,7 +30,7 @@
 
 setup(
     name='trendln',
-    version="0.1.17",
+    version="0.1.18",
     description='Support and Resistance Trend lines Calculator for Financial Analysis',
     long_description=long_description,
     long_description_content_type='text/markdown',

tests/test_trendln.py

@@ -431,6 +431,130 @@ def test_extend_to_end_date_passthrough(self):
         plt.close('all')
 
 
+# ---------------------------------------------------------------------------
+# Candlestick rendering via ohlc= parameter (issue #29)
+# ---------------------------------------------------------------------------
+
+def _make_ohlc(base):
+    """Build a simple synthetic OHLC 4-tuple from *base* close prices."""
+    opens  = [base[max(0, i - 1)] for i in range(len(base))]
+    highs  = [base[i] + 0.5 for i in range(len(base))]
+    lows   = [max(0.0, base[i] - 0.5) for i in range(len(base))]
+    closes = list(base)
+    return (opens, highs, lows, closes)
+
+
+class TestCandlestick:
+    """Tests for ohlc= parameter on plot_support_resistance / plot_sup_res_date."""
+
+    def test_ohlc_returns_figure(self):
+        import matplotlib
+        import matplotlib.pyplot as plt
+        ohlc = _make_ohlc(DATA_SIMPLE)
+        fig = plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                      ohlc=ohlc)
+        assert isinstance(fig, matplotlib.figure.Figure)
+        plt.close('all')
+
+    def test_ohlc_draws_one_patch_per_bar(self):
+        """One Rectangle patch should be added for every bar."""
+        import matplotlib.pyplot as plt
+        ohlc = _make_ohlc(DATA_SIMPLE)
+        fig = plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                      ohlc=ohlc)
+        ax = fig.axes[0]
+        assert len(ax.patches) == len(DATA_SIMPLE)
+        plt.close('all')
+
+    def test_ohlc_no_price_line_plotted(self):
+        """When ohlc is supplied, the plain price line (25-point line) should
+        not be drawn — only wicks (2-point each) and other trendlines."""
+        import matplotlib.pyplot as plt
+        ohlc = _make_ohlc(DATA_SIMPLE)
+        fig = plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                      ohlc=ohlc, show_average=False)
+        ax = fig.axes[0]
+        # No line with len(DATA_SIMPLE) x-data points should exist
+        full_series_lines = [ln for ln in ax.lines
+                             if len(ln.get_xdata()) == len(DATA_SIMPLE)]
+        assert len(full_series_lines) == 0
+        plt.close('all')
+
+    def test_no_ohlc_still_draws_price_line(self):
+        """Without ohlc the price line with len(DATA_SIMPLE) points IS drawn."""
+        import matplotlib.pyplot as plt
+        fig = plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                      show_average=False)
+        ax = fig.axes[0]
+        full_series_lines = [ln for ln in ax.lines
+                             if len(ln.get_xdata()) == len(DATA_SIMPLE)]
+        assert len(full_series_lines) == 1
+        plt.close('all')
+
+    def test_ohlc_patches_colors_bullish_bearish(self):
+        """Bullish bars (close >= open) should be green; bearish bars red."""
+        import matplotlib.pyplot as plt
+        # Build data where first bar is bullish (open < close)
+        # and second is bearish (open > close)
+        opens  = [1.0, 3.0] + [1.5] * (len(DATA_SIMPLE) - 2)
+        closes = [3.0, 1.0] + [1.5] * (len(DATA_SIMPLE) - 2)
+        highs  = [v + 0.5 for v in closes]
+        lows   = [v - 0.5 for v in closes]
+        ohlc = (opens, highs, lows, closes)
+        fig = plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                      ohlc=ohlc)
+        ax = fig.axes[0]
+        patches = ax.patches
+        assert patches[0].get_facecolor()[:3] == pytest.approx(
+            tuple(int(c, 16) / 255 for c in ['2c', 'a0', '2c']), abs=0.01)  # green
+        assert patches[1].get_facecolor()[:3] == pytest.approx(
+            tuple(int(c, 16) / 255 for c in ['d6', '27', '28']), abs=0.01)  # red
+        plt.close('all')
+
+    def test_ohlc_invalid_not_tuple_raises(self):
+        with pytest.raises(ValueError, match='ohlc'):
+            plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                    ohlc=[DATA_SIMPLE] * 4)
+
+    def test_ohlc_invalid_wrong_length_raises(self):
+        with pytest.raises(ValueError, match='ohlc'):
+            ohlc = _make_ohlc(DATA_SIMPLE)
+            plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                    ohlc=ohlc[:3])
+
+    def test_ohlc_mismatched_series_lengths_raises(self):
+        opens, highs, lows, closes = _make_ohlc(DATA_SIMPLE)
+        with pytest.raises(ValueError, match='same length'):
+            plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                    ohlc=(opens[:-1], highs, lows, closes))
+
+    def test_ohlc_forwarded_by_plot_sup_res_date(self):
+        """plot_sup_res_date passes ohlc through to plot_support_resistance."""
+        import matplotlib
+        import pandas as pd
+        import matplotlib.pyplot as plt
+        idx = pd.date_range('2020-01-01', periods=len(DATA_SIMPLE), freq='D')
+        ohlc = _make_ohlc(DATA_SIMPLE)
+        fig = plot_sup_res_date(DATA_SIMPLE, idx, extmethod=METHOD_NAIVE,
+                                ohlc=ohlc)
+        ax = fig.axes[0]
+        assert len(ax.patches) == len(DATA_SIMPLE)
+        assert isinstance(fig, matplotlib.figure.Figure)
+        plt.close('all')
+
+    def test_ohlc_compatible_with_ax_parameter(self):
+        """ohlc and ax can be used together (subplot embedding)."""
+        import matplotlib.pyplot as plt
+        fig, (ax1, ax2) = plt.subplots(1, 2)
+        ohlc = _make_ohlc(DATA_SIMPLE)
+        returned = plot_support_resistance(DATA_SIMPLE, extmethod=METHOD_NAIVE,
+                                           ohlc=ohlc, ax=ax2)
+        assert returned is fig
+        assert len(ax2.patches) == len(DATA_SIMPLE)
+        assert len(ax1.patches) == 0
+        plt.close('all')
+
+
 # ---------------------------------------------------------------------------
 # Input validation
 # ---------------------------------------------------------------------------

trendln/__init__.py

@@ -1050,26 +1050,63 @@ def _cluster(price_idx_pairs):
     resistance_levels = sorted(_cluster(resistance_pairs), key=lambda c: c[0])
     return support_levels, resistance_levels
 
+def _draw_candlesticks(ax, opens, highs, lows, closes, series_label):
+    """Draw OHLC candlestick bars on *ax* using matplotlib Rectangle patches.
+
+    Each bar has a thin wick spanning the full high-low range and a coloured
+    body from open to close (green = bullish, red = bearish).  The first bar
+    carries the legend label; subsequent bars are suppressed.
+    """
+    import matplotlib.patches as mpatches
+    opens  = [float(v) for v in opens]
+    highs  = [float(v) for v in highs]
+    lows   = [float(v) for v in lows]
+    closes = [float(v) for v in closes]
+    n = len(opens)
+    width = 0.6
+    bullish = '#2ca02c'
+    bearish = '#d62728'
+    lbl = series_label or 'OHLC'
+    for i in range(n):
+        color = bullish if closes[i] >= opens[i] else bearish
+        ax.plot([i, i], [lows[i], highs[i]], color='black', linewidth=0.8, zorder=2)
+        body_lo = min(opens[i], closes[i])
+        body_hi = max(opens[i], closes[i])
+        rect = mpatches.Rectangle(
+            (i - width / 2, body_lo), width, body_hi - body_lo,
+            facecolor=color, edgecolor='black', linewidth=0.5, zorder=3,
+            label=lbl)
+        ax.add_patch(rect)
+        lbl = '_nolegend_'
+
 def plot_sup_res_date(hist, idx, numbest = 2, fromwindows = True, pctbound=0.1,
                       extmethod = METHOD_NUMDIFF, method=METHOD_NSQUREDLOGN, window=125,
                       errpct = 0.005, hough_scale=0.01, hough_prob_iter=10, sortError=False, accuracy=2,
                       title='Prices with Support/Resistance Trend Lines', y_axis_label='Price', series_label=None,
-                      show_average=True, ax=None, extend_to_end=False):
+                      show_average=True, ax=None, extend_to_end=False, ohlc=None):
     import matplotlib.ticker as ticker
     return plot_support_resistance(hist, ticker.FuncFormatter(datefmt(idx)), numbest, fromwindows,
                                    pctbound, extmethod, method, window, errpct, hough_scale,
                                    hough_prob_iter, sortError, accuracy,
                                    title=title, y_axis_label=y_axis_label, series_label=series_label,
-                                   show_average=show_average, ax=ax, extend_to_end=extend_to_end)
+                                   show_average=show_average, ax=ax, extend_to_end=extend_to_end,
+                                   ohlc=ohlc)
 
 def plot_support_resistance(hist, xformatter = None, numbest = 2, fromwindows = True,
                             pctbound=0.1, extmethod = METHOD_NUMDIFF, method=METHOD_NSQUREDLOGN,
                             window=125, errpct = 0.005, hough_scale=0.01, hough_prob_iter=10, sortError=False, accuracy=2,
                             title='Prices with Support/Resistance Trend Lines', y_axis_label='Price', series_label=None,
-                            show_average=True, ax=None, extend_to_end=False):
+                            show_average=True, ax=None, extend_to_end=False, ohlc=None):
     import matplotlib.pyplot as plt
     import matplotlib.ticker as ticker
     ret = calc_support_resistance(hist, extmethod, method, window, errpct, hough_scale, hough_prob_iter, sortError, accuracy)
+    if ohlc is not None:
+        if (not isinstance(ohlc, tuple) or len(ohlc) != 4 or
+                not all(check_num_alike(s) for s in ohlc)):
+            raise ValueError(
+                'ohlc must be a 4-tuple of (opens, highs, lows, closes) numeric sequences')
+        if not (len(ohlc[0]) == len(ohlc[1]) == len(ohlc[2]) == len(ohlc[3])):
+            raise ValueError('all four ohlc sequences must have the same length')
     if ax is None:
         plt.clf()
         ax = plt.subplot(111)
@@ -1082,15 +1119,17 @@ def plot_support_resistance(hist, xformatter = None, numbest = 2, fromwindows =
             disp = [(hist[0], minimaIdxs, pmin, 'yo', 'Avg. Support', 'y--'), (hist[1], maximaIdxs, pmax, 'bo', 'Avg. Resistance', 'b--')]
             dispwin = [(hist[0], minwindows, 'Support', 'g--'), (hist[1], maxwindows, 'Resistance', 'r--')]
             disptrend = [(hist[0], mintrend, 'Support', 'g--'), (hist[1], maxtrend, 'Resistance', 'r--')]
-            ax.plot(range(len_h), hist[0], 'k--', label=f'Low {series_label or "Price"}')
-            ax.plot(range(len_h), hist[1], 'm--', label=f'High {series_label or "Price"}')
+            if ohlc is None:
+                ax.plot(range(len_h), hist[0], 'k--', label=f'Low {series_label or "Price"}')
+                ax.plot(range(len_h), hist[1], 'm--', label=f'High {series_label or "Price"}')
         else:
             len_h = len(hist)
             min_h, max_h = min(hist), max(hist)
             disp = [(hist, minimaIdxs, pmin, 'yo', 'Avg. Support', 'y--'), (hist, maximaIdxs, pmax, 'bo', 'Avg. Resistance', 'b--')]
             dispwin = [(hist, minwindows, 'Support', 'g--'), (hist, maxwindows, 'Resistance', 'r--')]
             disptrend = [(hist, mintrend, 'Support', 'g--'), (hist, maxtrend, 'Resistance', 'r--')]
-            ax.plot(range(len_h), hist, 'k--', label=series_label or 'Close Price')
+            if ohlc is None:
+                ax.plot(range(len_h), hist, 'k--', label=series_label or 'Close Price')
     else:
         minimaIdxs, pmin, mintrend, minwindows = ([], [], [], []) if hist[0] is None else ret
         maximaIdxs, pmax, maxtrend, maxwindows = ([], [], [], []) if hist[1] is None else ret
@@ -1099,7 +1138,12 @@ def plot_support_resistance(hist, xformatter = None, numbest = 2, fromwindows =
         disp = [(hist[1], maximaIdxs, pmax, 'bo', 'Avg. Resistance', 'b--') if hist[0] is None else (hist[0], minimaIdxs, pmin, 'yo', 'Avg. Support', 'y--')]
         dispwin = [(hist[1], maxwindows, 'Resistance', 'r--') if hist[0] is None else (hist[0], minwindows, 'Support', 'g--')]
         disptrend = [(hist[1], maxtrend, 'Resistance', 'r--') if hist[0] is None else (hist[0], mintrend, 'Support', 'g--')]
-        ax.plot(range(len_h), hist[1 if hist[0] is None else 0], 'k--', label= ('High' if hist[0] is None else 'Low') + ' Price')
+        if ohlc is None:
+            ax.plot(range(len_h), hist[1 if hist[0] is None else 0], 'k--', label= ('High' if hist[0] is None else 'Low') + ' Price')
+    if ohlc is not None:
+        min_h = min(float(v) for v in ohlc[2])  # lows
+        max_h = max(float(v) for v in ohlc[1])  # highs
+        _draw_candlesticks(ax, ohlc[0], ohlc[1], ohlc[2], ohlc[3], series_label)
     for h, idxs, pm, clrp, lbl, clrl in disp:
         ax.plot(idxs, [h[x] for x in idxs], clrp)
         if show_average: