Handle the case where the passed in str is a table itself

yuanxu-li · yuanxu-li · commit d1b3ac2aafa1 · 2019-02-23T15:53:50.000-08:00
diff --git a/html_table_extractor/extractor.py b/html_table_extractor/extractor.py
@@ -6,19 +6,22 @@
 import csv
 import pdb
 
-
 class Extractor(object):
-    def __init__(self, table, id_=None, **kwargs):
-        # input is Tag
-        if isinstance(table, Tag):
-            self._table = table.find(id=id_)
-        # input is str
-        # for py2, make sure you pass in str
-        # for py3, everything is str by default
-        elif isinstance(table, str):
-            self._table = BeautifulSoup(table, 'html.parser').find(id=id_)
+    def __init__(self, input, id_=None, **kwargs):
+        # TODO: should divide this class into two subclasses
+        # to deal with string and bs4.Tag separately
+
+        # validate the input
+        if not isinstance(input, str) and not isinstance(input, Tag):
+            raise Exception('Unrecognized type. Valid input: str, bs4.element.Tag')
+
+        soup = BeautifulSoup(input, 'html.parser').find() if isinstance(input, str) else input
+
+        # locate the target table
+        if soup.name == 'table':
+            self._table = soup
         else:
-            raise Exception('unrecognized type')
+            self._table = soup.find(id=id_)
 
         if 'transformer' in kwargs:
             self._transformer = kwargs['transformer']
diff --git a/tests/tests.py b/tests/tests.py
@@ -63,7 +63,8 @@ def test_config_transformer(self):
 class TestPassId(unittest.TestCase):
     def test_init_with_id(self):
         html = """
-        <table id='wanted'>
+        <body>
+          <table id='wanted'>
             <tr>
               <td>1</td>
               <td>2</td>
@@ -72,12 +73,13 @@ def test_init_with_id(self):
               <td>3</td>
               <td>4</td>
             </tr>
-        </table>
-        <table id='unwanted'>
+          </table>
+          <table id='unwanted'>
             <tr>
               <td>unwanted</td>
             </tr>
-        </table>
+          </table>
+        </body>
         """
         soup = BeautifulSoup(html, 'html.parser')
         extractor = Extractor(soup, id_='wanted').parse()
@@ -86,6 +88,7 @@ def test_init_with_id(self):
             [[u'1', u'2'], [u'3', u'4']]
         )
 
+
 class TestComplexExtractor(unittest.TestCase):
     def setUp(self):
         html = """
