diff --git a/README.md b/README.md index 7c7cc6d..149cc8e 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![codecov](https://codecov.io/gh/jealous/stockstats/branch/master/graph/badge.svg?token=IFMD1pVJ7T)](https://codecov.io/gh/jealous/stockstats) [![pypi](https://img.shields.io/pypi/v/stockstats.svg)](https://pypi.python.org/pypi/stockstats) -VERSION: 0.5.1 +VERSION: 0.5.2 ## Introduction @@ -538,6 +538,10 @@ The Bollinger bands includes three lines * `df['boll_ub']` is the upper band * `df['boll_lb']` is the lower band +The default window of boll is defined by `BOLL_PERIOD`. The default value is 20. +You can also supply your window with `df['boll_10']`. It will also +generate the `boll_ub_10` and `boll_lb_10` column. + The default period of the Bollinger Band can be changed with `StockDataFrame.BOLL_PERIOD`. The width of the bands can be turned with `StockDataFrame.BOLL_STD_TIMES`. The default value is 2. diff --git a/stockstats.py b/stockstats.py index eda718d..37767b6 100644 --- a/stockstats.py +++ b/stockstats.py @@ -916,7 +916,7 @@ def _get_ema(self, column, windows): column_name = '{}_{}_ema'.format(column, window) self[column_name] = self._ema(self[column], window) - def _get_boll(self): + def _get_boll(self, window=None): """ Get Bollinger bands. boll_ub means the upper band of the Bollinger bands @@ -927,12 +927,23 @@ def _get_boll(self): K = BOLL_STD_TIMES :return: None """ - moving_avg = self._sma(self['close'], self.BOLL_PERIOD) - moving_std = self._mstd(self['close'], self.BOLL_PERIOD) - self['boll'] = moving_avg + if window is None: + n = self.BOLL_PERIOD + boll = 'boll' + boll_ub = 'boll_ub' + boll_lb = 'boll_lb' + else: + n = self.get_int_positive(window) + boll = 'boll_{}'.format(n) + boll_ub = 'boll_ub_{}'.format(n) + boll_lb = 'boll_lb_{}'.format(n) + moving_avg = self._sma(self['close'], n) + moving_std = self._mstd(self['close'], n) + + self[boll] = moving_avg width = self.BOLL_STD_TIMES * moving_std - self['boll_ub'] = moving_avg + width - self['boll_lb'] = moving_avg - width + self[boll_ub] = moving_avg + width + self[boll_lb] = moving_avg - width def _get_macd(self): """ Moving Average Convergence Divergence diff --git a/test.py b/test.py index e7d7b21..33b0d18 100644 --- a/test.py +++ b/test.py @@ -323,6 +323,16 @@ def test_bollinger(self): assert_that(boll_ub.loc[idx], near_to(10.1310)) assert_that(boll_lb.loc[idx], near_to(9.4759)) + def test_bollinger_with_window(self): + stock = self.get_stock().within(20140930, 20141211) + _ = stock['boll_20'] + idx = 20141103 + assert_that(stock['boll_20'].loc[idx], near_to(9.8035)) + assert_that(stock['boll_ub_20'].loc[idx], near_to(10.1310)) + assert_that(stock['boll_lb_20'].loc[idx], near_to(9.4759)) + _ = stock['boll_5'] + assert_that(stock['boll_ub_5'].loc[idx], near_to(10.44107)) + def test_bollinger_empty(self): stock = self.get_stock().within(18800101, 18900101) s = stock['boll_ub']