20
20
21
21
from six .moves import zip
22
22
23
- def _check_input (opens , closes , highs , lows ):
23
+ def _check_input (opens , closes , highs , lows , colors = None ):
24
24
"""Checks that *opens*, *highs*, *lows* and *closes* have the same length.
25
25
NOTE: this code assumes if any value open, high, low, close is
26
26
missing (*-1*) they all are missing
@@ -46,6 +46,10 @@ def _check_input(opens, closes, highs, lows):
46
46
if not same_length :
47
47
raise ValueError ('O,H,L,C must have the same length!' )
48
48
49
+ if colors :
50
+ if len (opens ) != len (colors ):
51
+ raise ValueError ('O,H,L,C and Colors must have the same length!' )
52
+
49
53
o = np .where (np .isnan (opens ))[0 ]
50
54
h = np .where (np .isnan (highs ))[0 ]
51
55
l = np .where (np .isnan (lows ))[0 ]
@@ -85,11 +89,11 @@ def _check_and_convert_xlim_configuration(data, config):
85
89
return xlim
86
90
87
91
88
- def _construct_mpf_collections (ptype ,dates ,xdates ,opens ,highs ,lows ,closes ,volumes ,config ,style ):
92
+ def _construct_mpf_collections (ptype ,dates ,xdates ,opens ,highs ,lows ,closes ,volumes ,config ,style , colors ):
89
93
collections = None
90
94
if ptype == 'candle' or ptype == 'candlestick' :
91
95
collections = _construct_candlestick_collections (xdates , opens , highs , lows , closes ,
92
- marketcolors = style ['marketcolors' ],config = config )
96
+ marketcolors = style ['marketcolors' ],config = config , colors = colors )
93
97
94
98
elif ptype == 'hollow_and_filled' :
95
99
collections = _construct_hollow_candlestick_collections (xdates , opens , highs , lows , closes ,
@@ -176,16 +180,45 @@ def coalesce_volume_dates(in_volumes, in_dates, indexes):
176
180
return volumes , dates
177
181
178
182
179
- def _updown_colors (upcolor ,downcolor ,opens ,closes ,use_prev_close = False ):
180
- if upcolor == downcolor :
181
- return upcolor
182
- cmap = {True : upcolor , False : downcolor }
183
- if not use_prev_close :
184
- return [ cmap [opn < cls ] for opn ,cls in zip (opens ,closes ) ]
183
+ def _updown_colors (upcolor ,downcolor ,opens ,closes ,use_prev_close = False ,colors = None ):
184
+ if not colors :
185
+ if upcolor == downcolor :
186
+ return upcolor
187
+ cmap = {True : upcolor , False : downcolor }
188
+ if not use_prev_close :
189
+ return [ cmap [opn < cls ] for opn ,cls in zip (opens ,closes ) ]
190
+ else :
191
+ first = cmap [opens [0 ] < closes [0 ]]
192
+ _list = [ cmap [pre < cls ] for cls ,pre in zip (closes [1 :], closes ) ]
193
+ return [first ] + _list
185
194
else :
186
- first = cmap [opens [0 ] < closes [0 ]]
187
- _list = [ cmap [pre < cls ] for cls ,pre in zip (closes [1 :], closes ) ]
188
- return [first ] + _list
195
+ cmap = {True : 'up' , False : 'down' }
196
+ default = {'up' : upcolor , 'down' : downcolor }
197
+ custom = []
198
+ if not use_prev_close :
199
+ for i in range (len (opens )):
200
+ opn = opens [i ]
201
+ cls = closes [i ]
202
+ if colors [i ]:
203
+ custom .append (colors [i ][cmap [opn < cls ]])
204
+ else :
205
+ custom .append (default [cmap [opn < cls ]])
206
+ else :
207
+ if color [0 ]:
208
+ custom .append (colors [0 ][cmap [opens [0 ] < closes [0 ]]])
209
+ else :
210
+ custom .append (default [cmap [opens [0 ] < closes [0 ]]])
211
+
212
+ for i in range (len (closes ) - 1 ):
213
+ pre = closes [1 :][i ]
214
+ cls = closes [i ]
215
+ if colors [i ]:
216
+ custom .append (colors [i ][cmap [pre < cls ]])
217
+ else :
218
+ custom .append (default [cmap [pre < cls ]])
219
+
220
+ return custom
221
+
189
222
190
223
191
224
def _updownhollow_colors (upcolor ,downcolor ,hollowcolor ,opens ,closes ):
@@ -525,7 +558,7 @@ def _construct_ohlc_collections(dates, opens, highs, lows, closes, marketcolors=
525
558
return [rangeCollection , openCollection , closeCollection ]
526
559
527
560
528
- def _construct_candlestick_collections (dates , opens , highs , lows , closes , marketcolors = None , config = None ):
561
+ def _construct_candlestick_collections (dates , opens , highs , lows , closes , marketcolors = None , config = None , colors = None ):
529
562
"""Represent the open, close as a bar line and high low range as a
530
563
vertical line.
531
564
@@ -552,8 +585,8 @@ def _construct_candlestick_collections(dates, opens, highs, lows, closes, market
552
585
ret : list
553
586
(lineCollection, barCollection)
554
587
"""
555
-
556
- _check_input (opens , highs , lows , closes )
588
+
589
+ _check_input (opens , highs , lows , closes , colors )
557
590
558
591
if marketcolors is None :
559
592
marketcolors = _get_mpfstyle ('classic' )['marketcolors' ]
@@ -581,17 +614,34 @@ def _construct_candlestick_collections(dates, opens, highs, lows, closes, market
581
614
582
615
alpha = marketcolors ['alpha' ]
583
616
617
+ candle_c = None
618
+ wick_c = None
619
+ edge_c = None
620
+ if colors :
621
+ candle_c = []
622
+ wick_c = []
623
+ edge_c = []
624
+ for color in colors :
625
+ if color :
626
+ candle_c .append ({'up' : mcolors .to_rgba (color ['candle' ]['up' ], alpha ), 'down' : mcolors .to_rgba (color ['candle' ]['down' ], alpha )})
627
+ wick_c .append ({'up' : mcolors .to_rgba (color ['wick' ]['up' ], 1 ), 'down' : mcolors .to_rgba (color ['wick' ]['down' ], 1 )})
628
+ edge_c .append ({'up' : mcolors .to_rgba (color ['edge' ]['up' ], 1 ), 'down' : mcolors .to_rgba (color ['edge' ]['down' ], 1 )})
629
+ else :
630
+ candle_c .append (None )
631
+ wick_c .append (None )
632
+ edge_c .append (None )
633
+
584
634
uc = mcolors .to_rgba (marketcolors ['candle' ][ 'up' ], alpha )
585
635
dc = mcolors .to_rgba (marketcolors ['candle' ]['down' ], alpha )
586
- colors = _updown_colors (uc , dc , opens , closes )
636
+ colors = _updown_colors (uc , dc , opens , closes , colors = candle_c )
587
637
588
638
uc = mcolors .to_rgba (marketcolors ['edge' ][ 'up' ], 1.0 )
589
639
dc = mcolors .to_rgba (marketcolors ['edge' ]['down' ], 1.0 )
590
- edgecolor = _updown_colors (uc , dc , opens , closes )
640
+ edgecolor = _updown_colors (uc , dc , opens , closes , colors = edge_c )
591
641
592
642
uc = mcolors .to_rgba (marketcolors ['wick' ][ 'up' ], 1.0 )
593
643
dc = mcolors .to_rgba (marketcolors ['wick' ]['down' ], 1.0 )
594
- wickcolor = _updown_colors (uc , dc , opens , closes )
644
+ wickcolor = _updown_colors (uc , dc , opens , closes , colors = wick_c )
595
645
596
646
lw = config ['_width_config' ]['candle_linewidth' ]
597
647
0 commit comments