#! usr/bin/env python3
from docopt import docopt
from matplotlib import pyplot
import re

"""SQLViz

Usage:
    sqlviz [-hnkdlpo DIR] FILE

Options:
    -h --help       Display this help file
    -k --keys       Graph the number of foreign keys vs primary keys
    -d --datatypes  Graph the distribution of datatypes
    -l --lengths    Graph the distribution of lengths for each data type
    -p --print      Print text to the console as well as creating graphs
    -n --no-display Don't display the generated graphs.
    -o DIR          Output graphs to the specified directory

"""

class Schema:
    """
    Wraps the SQL source code for a schema and provides methods to get information about that schema.
    """

    table_def = re.compile(r"CREATE TABLE|create table")
    primary_key = re.compile(r"PRIMARY KEY|primary key")
    foreign_key = re.compile(r"FOREIGN KEY|foreign key")
    varchar = re.compile(r"(?:VARCHAR|varchar)\s*\((\d+)\)")
    decimal = re.compile(r"(?:DECIMAL|decimal)\s*\((\d+,\d+)\)")
    decimal_extract = re.compile(r"(?P<p>\d+),\s*(?P<d>\d+)")
    integer = re.compile(r"(INT|int|INTEGER|integer)(\s|,)")
    text = re.compile(r"(TEXT|text)(\s|,)")
    numeric = re.compile(r"(?:NUMERIC|numeric)\s*\((\d+,\s*\d+)\)")

    def __init__(self, source):
        """
        Creates a new instance of Schema for the specified source code string.
        """
        self.source = source

    def n_tables(self):
        """
        Returns the number of tables defined in the schema
        """
        return len(Schema.table_def.findall(self.source))

    def n_keys(self):
        """
        Returns the number of keys defined in the schema
        """
        return {"PRIMARY KEY": len(Schema.primary_key.findall(self.source)),
                "FOREIGN KEY": len(Schema.foreign_key.findall(self.source))}

    def n_datatypes(self):
        """
        Returns the number of each data type in the schema.
        """
        return {"INT": len(Schema.integer.findall(self.source)),
                "DECIMAL": len(Schema.decimal.findall(self.source)),
                "NUMERIC": len(Schema.numeric.findall(self.source)),
                "VARCHAR": len(Schema.varchar.findall(self.source)),
                "TEXT": len(Schema.text.findall(self.source))}

    def lengths(self):
        """
        Returns a dictionary mapping each data type in the schema
        to a list of the lengths of those data types.
        """
        return {"VARCHAR": [int(v) for v in Schema.varchar.findall(self.source)],
                "DECIMAL": [(int(d.group("p")), int(d.group("d"))) for d in map(Schema.decimal_extract.search,Schema.decimal.findall(self.source))],
                "NUMERIC": [(int(d.group("p")), int(d.group("d"))) for d in map(Schema.decimal_extract.search,Schema.numeric.findall(self.source))]
        }

if __name__ == "__main__":
    opts = docopt(__doc__, help=True, version="0.1")

@@ -2918,8 +2918,762 @@
 =%220.1%22)%0A
+%0A    with open(opts%5B%22FILE%22%5D, 'r') as f:%0A        source = f.read()%0A%0A    schema = schema(source)%0A%0A    # Begin plotting%0A    fignum = 0%0A%0A    if opts%5B%22--keys%22%5D: # pie chart of keys%0A        fignum = fignum + 1%0A        figure(fignum, figsize=(6,6))%0A        ax = axes(%5B0.1, 0.1, 0.8, 0.8%5D)%0A%0A        keys = schema.n_keys()%0A        total_keys = keys%5B%22PRIMARY KEY%22%5D + keys%5B%22FOREIGN KEY%22%5D%0A%0A        fracs = %5B # determine fractions of primary/foreign%0A            (keys%5B%22PRIMARY_KEY%22%5D/total_keys)*100, (keys%5B%22FOREIGN KEY%22%5D/total_keys)*100,%0A            %5D%0A%0A        pie(fracs, labels = %5B%22primary%22, %22foreign%22%5D)%0A        title(%22Key Composition%22)%0A%0A        if not opts%5B%22--no-display%22%5D:%0A            show()%0A        if opts%5B%22-o%22%5D:%0A            savefig((opts%5B%22-o%22%5D + %22/keys.pdf%22))%0A%0A

Add pie charts for key comp to sqlviz