Merge branch 'master' of https://github.com/dstephens99/pandas into dstephens99-master

Author: jreback

pandas/io/data.py

@@ -736,9 +736,8 @@ def _get_option_data(self, month, year, expiry, name):
                              " found".format(table_loc, ntables))
 
         option_data = _parse_options_data(tables[table_loc])
-        option_data = self._process_data(option_data)
         option_data['Type'] = name[:-1]
-        option_data.set_index(['Strike', 'Expiry', 'Type', 'Symbol'], inplace=True)
+        option_data = self._process_data(option_data, name[:-1])
 
         if month == CUR_MONTH and year == CUR_YEAR:
             setattr(self, name, option_data)
@@ -859,8 +858,7 @@ def get_near_stock_price(self, above_below=2, call=True, put=False,
                              month=None, year=None, expiry=None):
         """
         ***Experimental***
-        Cuts the data frame opt_df that is passed in to only take
-        options that are near the current stock price.
+        Returns a data frame of options that are near the current stock price.
 
         Parameters
         ----------
@@ -889,34 +887,38 @@ def get_near_stock_price(self, above_below=2, call=True, put=False,
          Note: Format of returned data frame is dependent on Yahoo and may change.
 
         """
-        year, month, expiry = self._try_parse_dates(year, month, expiry)
 
         to_ret = Series({'calls': call, 'puts': put})
         to_ret = to_ret[to_ret].index
 
         data = {}
 
         for nam in to_ret:
-            if month:
-                m1 = _two_char_month(month)
-                name = nam + m1 + str(year)[2:]
+            df = self._get_option_data(month, year, expiry, nam)
+            data[nam] = self.chop_data(df, above_below, self.underlying_price)
+
+        return concat([data[nam] for nam in to_ret]).sortlevel()
+
+    def chop_data(self, df, above_below=2, underlying_price=None):
+        """Returns a data frame only options that are near the current stock price."""
 
+        if not underlying_price:
             try:
-                df = getattr(self, name)
+                underlying_price = self.underlying_price
             except AttributeError:
-                meth_name = 'get_{0}_data'.format(nam[:-1])
-                df = getattr(self, meth_name)(expiry=expiry)
+                underlying_price = np.nan
 
-            if self.underlying_price:
-                start_index = np.where(df.index.get_level_values('Strike')
-                                   > self.underlying_price)[0][0]
+        if not np.isnan(underlying_price):
+            start_index = np.where(df.index.get_level_values('Strike')
+                                   > underlying_price)[0][0]
 
-                get_range = slice(start_index - above_below,
+            get_range = slice(start_index - above_below,
                               start_index + above_below + 1)
-                chop = df[get_range].dropna(how='all')
-                data[nam] = chop
+            df = df[get_range].dropna(how='all')
+
+        return df
+
 
-        return concat([data[nam] for nam in to_ret]).sortlevel()
 
     @staticmethod
     def _try_parse_dates(year, month, expiry):
@@ -1048,7 +1050,7 @@ def get_forward_data(self, months, call=True, put=False, near=False,
                     frame = self.get_near_stock_price(call=call, put=put,
                                                       above_below=above_below,
                                                       month=m2, year=y2)
-                frame = self._process_data(frame)
+                frame = self._process_data(frame, name[:-1])
 
                 all_data.append(frame)
 
@@ -1178,7 +1180,7 @@ def _parse_url(self, url):
         return root
 
 
-    def _process_data(self, frame):
+    def _process_data(self, frame, type):
         """
         Adds columns for Expiry, IsNonstandard (ie: deliverable is not 100 shares)
         and Tag (the tag indicating what is actually deliverable, None if standard).
@@ -1195,5 +1197,7 @@ def _process_data(self, frame):
         frame['Underlying_Price'] = self.underlying_price
         frame["Quote_Time"] = self.quote_time
         frame.rename(columns={'Open Int': 'Open_Int'}, inplace=True)
+        frame['Type'] = type
+        frame.set_index(['Strike', 'Expiry', 'Type', 'Symbol'], inplace=True)
 
         return frame

pandas/io/tests/test_data.py

@@ -250,6 +250,9 @@ def setUpClass(cls):
         cls.html2 = os.path.join(cls.dirpath, 'yahoo_options2.html')
         cls.root1 = cls.aapl._parse_url(cls.html1)
         cls.root2 = cls.aapl._parse_url(cls.html2)
+        cls.tables1 = cls.aapl._parse_option_page_from_yahoo(cls.root1)
+        cls.unprocessed_data1 = web._parse_options_data(cls.tables1[cls.aapl._TABLE_LOC['puts']])
+        cls.data1 = cls.aapl._process_data(cls.unprocessed_data1, 'put')
 
     @classmethod
     def tearDownClass(cls):
@@ -324,6 +327,13 @@ def test_sample_page_price_quote_time1(self):
         self.assertIsInstance(price, (int, float, complex))
         self.assertIsInstance(quote_time, (datetime, Timestamp))
 
+    def test_chop(self):
+        #regression test for #7625
+        self.aapl.chop_data(self.data1, above_below=2, underlying_price=np.nan)
+        chopped = self.aapl.chop_data(self.data1, above_below=2, underlying_price=300)
+        self.assertIsInstance(chopped, DataFrame)
+        self.assertTrue(len(chopped) > 1)
+
     @network
     def test_sample_page_price_quote_time2(self):
         #Tests the weekday quote time format
@@ -334,10 +344,7 @@ def test_sample_page_price_quote_time2(self):
     @network
     def test_sample_page_chg_float(self):
         #Tests that numeric columns with comma's are appropriately dealt with
-        tables = self.aapl._parse_option_page_from_yahoo(self.root1)
-        data = web._parse_options_data(tables[self.aapl._TABLE_LOC['puts']])
-        option_data = self.aapl._process_data(data)
-        self.assertEqual(option_data['Chg'].dtype, 'float64')
+        self.assertEqual(self.data1['Chg'].dtype, 'float64')
 
 
 class TestOptionsWarnings(tm.TestCase):