Merge pull request #8 from alisaraa/aa-implement-a-sqlite-db

Added SQLite support

Author: mariostoev

finviz/save_data.py

@@ -1,5 +1,17 @@
 import csv
+import os
+import sqlite3
+import re
 
+def create_connection():
+    
+    sqlite_file = "../screener.sqlite"
+    
+    try:
+        conn = sqlite3.connect(sqlite_file)
+        return conn
+    except sqlite3.Error as error:
+        raise ("An error has occurred while connecting to the database: ", error.args[0])
 
 def export_to_csv(headers, data, directory):
 
@@ -9,3 +21,58 @@ def export_to_csv(headers, data, directory):
 
         for n in data:
             dict_writer.writerows(n)
+
+def export_to_db(headers, data):
+
+    field_list = ""
+    table_name = "screener_results"  # name of the table to be created
+    conn = create_connection()
+    c = conn.cursor()
+
+    for field in headers:
+
+        field_cleaned = re.sub(r'[^\w\s]', '', field)
+        field_cleaned = field_cleaned.replace(" ", "")
+        field_list += field_cleaned + " TEXT, "
+
+    c.execute("CREATE TABLE IF NOT EXISTS {tn} ({fl})"
+              .format(tn=table_name, fl=field_list[:-2]))
+
+    inserts = ""
+    for data in data:
+
+        for level in data:
+
+            insert_fields = "("
+            for field, value in level.items():
+                insert_fields += "\"" + value + "\", "
+
+            inserts += insert_fields[:-2] + "), "
+
+    insert_lines = inserts[:-2]
+
+    try:
+        c.execute("INSERT INTO {tn} VALUES {iv}".
+                  format(tn=table_name, iv=insert_lines))
+    except sqlite3.Error as error:
+        print("An error has occurred", error.args[0])
+
+    conn.commit()
+    conn.close()
+
+
+def select_from_db():
+
+    conn = create_connection()
+    c = conn.cursor()
+    c.execute("SELECT * FROM screener_results")
+
+    rows = c.fetchall()
+
+    for row in rows:
+        print(row)
+
+    rows = c.fetchall()
+
+    for row in rows:
+        print(row)

finviz/screener.py

@@ -1,12 +1,13 @@
+from .save_data import export_to_db, select_from_db, export_to_csv
 from lxml import html
 from lxml import etree
 import finviz.request_functions as send
 import finviz.scraper_functions as scrape
-
+from os import getcwd
 
 class Screener(object):
 
-    def __init__(self, tickers=None, filters=None, rows=None, order='', signal='', table='Overview'):
+    def __init__(self, tickers=None, filters=None, order='', rows=None, signal='', table='Overview'):
 
         if tickers is None:
             self.tickers = []
@@ -32,15 +33,42 @@ def __init__(self, tickers=None, filters=None, rows=None, order='', signal='', t
 
     def to_csv(self, directory=None):
 
-        from .save_data import export_to_csv
-
         if directory is None:
-
-            import os
-            directory = os.getcwd()
+            directory = getcwd()
 
         export_to_csv(self.headers, self.data, directory)
 
+    def to_sqlite(self):
+        export_to_db(self.headers, self.data)
+
+    def display_db(self):
+        select_from_db()
+
+    def __get_total_rows(self):
+
+        total_element = self.page_content.cssselect('td[width="140"]')
+        self.rows = int(etree.tostring(total_element[0]).decode("utf-8").split('</b>')[1].split(' ')[0])
+
+    def __get_page_urls(self):
+
+        try:
+            total_pages = int([i.text.split('/')[1] for i in self.page_content.cssselect('option[value="1"]')][0])
+        except IndexError:  # No results found
+            return None
+
+        urls = []
+
+        for page_number in range(1, total_pages + 1):
+
+            sequence = 1 + (page_number - 1) * 20
+
+            if sequence - 20 <= self.rows < sequence:
+                break
+            else:
+                urls.append(self.url + '&r={}'.format(str(sequence)))
+
+        self.page_urls = urls
+
     def __get_table_headers(self):
 
         first_row = self.page_content.cssselect('tr[valign="middle"]')