Added documentation for classes and methods

Author: mariostoev

finviz/error_handling.py

@@ -0,0 +1,14 @@
+class NoResults(Exception):
+    """ Raise when there are no results found. """
+
+    def __init__(self, query):
+        self.query = query
+        super(NoResults, self).__init__("No results found for query: {}".format(query))
+
+
+class InvalidTableType(Exception):
+    """ Raise when the given table type is invalid. """
+
+    def __init__(self, arg):
+        self.arg = arg
+        super(InvalidTableType, self).__init__("Invalid table type called: {}".format(arg))

finviz/scraper_functions.py

@@ -6,8 +6,12 @@ def get_total_rows(page_content):
     """ Returns the total number of rows(results). """
 
     total_element = page_content.cssselect('td[width="140"]')
+    total_number = etree.tostring(total_element[0]).decode('utf-8').split('</b>')[1].split()[0]
 
-    return int(etree.tostring(total_element[0]).decode("utf-8").split('</b>')[1].split(' ')[0])
+    try:
+        return int(total_number)
+    except ValueError:
+        return 0
 
 
 def get_page_urls(page_content, rows, url):

finviz/screener.py

@@ -1,14 +1,14 @@
 from finviz.request_functions import Connector, http_request
+from finviz.error_handling import NoResults, InvalidTableType
 from .save_data import export_to_db, export_to_csv
 from urllib.parse import urlencode
 from lxml import html
 from lxml import etree
 import finviz.scraper_functions as scrape
 
-# TODO > Add __add__ and __slice__(?) methods to the Screener class
-# TODO > Create Error class
 # TODO > Add unittests
-# TODO > Improve performance
+# TODO > Make self.data list of dicts, not lists of lists of dicts
+# TODO > Implement __add__, __slice__, __iter__, __getitem__
 
 
 class Screener(object):
@@ -56,23 +56,30 @@ def __init__(self, tickers=None, filters=None, rows=None, order='', signal='', t
 
         self._order = order
         self._signal = signal
-        self._table = self._table_types[table]
+
+        try:
+            self._table = self._table_types[table]
+        except KeyError:
+            raise InvalidTableType(table)
+
         self._page_unparsed, self._url = http_request('https://finviz.com/screener.ashx', payload={
                                                    'v': self._table,
                                                    't': ','.join(self._tickers),
                                                    'f': ','.join(self._filters),
                                                    'o': self._order,
                                                    's': self._signal
                                                    })
-
         self._page_content = html.fromstring(self._page_unparsed)
-        self._headers = self.__get_table_headers()
+        self._total_rows = scrape.get_total_rows(self._page_content)
 
-        if rows is None:
-            self._rows = scrape.get_total_rows(self._page_content)
+        if self._total_rows == 0:
+            raise NoResults(self._url.split('?')[1])
+        elif rows is None or rows > self._total_rows:
+            self._rows = self._total_rows
         else:
             self._rows = rows
 
+        self._headers = self.__get_table_headers()
         self.data = self.__search_screener()
 
     def __repr__(self):
@@ -114,6 +121,9 @@ def __len__(self):
 
         return int(self._rows)
 
+    def __getitem__(self, position):
+        return self.data
+
     def to_sqlite(self):
         """ Exports the generated table into a SQLite database, located in the user's current directory. """
 
@@ -192,7 +202,7 @@ def scrape_row(line):
 
         for row in all_rows:
 
-            if int(row[0].text) is self._rows:
+            if int(row[0].text) == self._rows:
                 values = dict(zip(self._headers, scrape_row(row)))
                 data_sets.append(values)
                 break