Allow column names to be customized

Author: borgstrom

src/mplfinance/_arg_validators.py

@@ -2,7 +2,7 @@
 import pandas as pd
 import numpy  as np
 
-def _check_and_prepare_data(data):
+def _check_and_prepare_data(data, config):
     '''
     Check and Prepare the data input:
     For now, data must be a Pandas DataFrame with a DatetimeIndex
@@ -25,16 +25,17 @@ def _check_and_prepare_data(data):
     if not isinstance(data.index,pd.core.indexes.datetimes.DatetimeIndex):
         raise TypeError('Expect data.index as DatetimeIndex')
 
-    cols = ['Open','High','Low','Close']
+    o, h, l, c, v = config["columns"]
+    cols = [o, h, l, c]
 
     dates   = mdates.date2num(data.index.to_pydatetime())
-    opens   = data['Open'].values
-    highs   = data['High'].values
-    lows    = data['Low'].values
-    closes  = data['Close'].values
-    if 'Volume' in data.columns:
-        volumes = data['Volume'].values
-        cols.append('Volume')
+    opens   = data[o].values
+    highs   = data[h].values
+    lows    = data[l].values
+    closes  = data[c].values
+    if v in data.columns:
+        volumes = data[v].values
+        cols.append(v)
     else:
         volumes = None
 

src/mplfinance/plotting.py

@@ -64,6 +64,10 @@ def _valid_plot_kwargs():
     '''
 
     vkwargs = {
+        'columns'     : { 'Default'     : ('Open', 'High', 'Low', 'Close', 'Volume'),
+                          'Validator'   : lambda value: isinstance(value, (tuple, list))
+                                                        and len(value) == 5
+                                                        and all(isinstance(c, str) for c in value) },
         'type'        : { 'Default'     : 'ohlc',
                           'Validator'   : lambda value: value in ['candle','candlestick','ohlc','bars','ohlc_bars','line', 'renko'] },
 
@@ -164,10 +168,10 @@ def plot( data, **kwargs ):
     Also provide ability to plot trading signals, and/or addtional user-defined data.
     """
 
-    dates,opens,highs,lows,closes,volumes = _check_and_prepare_data(data)
-
     config = _process_kwargs(kwargs, _valid_plot_kwargs())
 
+    dates,opens,highs,lows,closes,volumes = _check_and_prepare_data(data, config)
+
     if config['type'] == 'renko' and config['addplot'] is not None:
         err = "`addplot` is not supported for `type='renko'`"
         raise ValueError(err)