databases.py

users = [[0, "Hero", 0],
         [1, "Dunn", 2],
         [2, "Sue", 3],
         [3, "Chi", 3]]

from typing import Tuple, Sequence, List, Any, Callable, Dict, Iterator
from collections import defaultdict

# A few type aliases we'll use later
Row = Dict[str, Any]                        # A database row
WhereClause = Callable[[Row], bool]         # Predicate for a single row
HavingClause = Callable[[List[Row]], bool]  # Predicate over multiple rows

class Table:
    def __init__(self, columns: List[str], types: List[type]) -> None:
        assert len(columns) == len(types), "# of columns must == # of types"

        self.columns = columns         # Names of columns
        self.types = types             # Data types of columns
        self.rows: List[Row] = []      # (no data yet)

    def col2type(self, col: str) -> type:
        idx = self.columns.index(col)      # Find the index of the column,
        return self.types[idx]             # and return its type.

    def insert(self, values: list) -> None:
        # Check for right # of values
        if len(values) != len(self.types):
            raise ValueError(f"You need to provide {len(self.types)} values")

        # Check for right types of values
        for value, typ3 in zip(values, self.types):
            if not isinstance(value, typ3) and value is not None:
                raise TypeError(f"Expected type {typ3} but got {value}")

        # Add the corresponding dict as a "row"
        self.rows.append(dict(zip(self.columns, values)))

    def __getitem__(self, idx: int) -> Row:
        return self.rows[idx]

    def __iter__(self) -> Iterator[Row]:
        return iter(self.rows)

    def __len__(self) -> int:
        return len(self.rows)

    def __repr__(self):
        """Pretty representation of the table: columns then rows"""
        rows = "\n".join(str(row) for row in self.rows)

        return f"{self.columns}\n{rows}"

    def update(self,
               updates: Dict[str, Any],
               predicate: WhereClause = lambda row: True):
        # First make sure the updates have valid names and types
        for column, new_value in updates.items():
            if column not in self.columns:
                raise ValueError(f"invalid column: {column}")

            typ3 = self.col2type(column)
            if not isinstance(new_value, typ3) and new_value is not None:
                raise TypeError(f"expected type {typ3}, but got {new_value}")

        # Now update
        for row in self.rows:
            if predicate(row):
                for column, new_value in updates.items():
                    row[column] = new_value

    def delete(self, predicate: WhereClause = lambda row: True) -> None:
        """Delete all rows matching predicate"""
        self.rows = [row for row in self.rows if not predicate(row)]

    def select(self,
               keep_columns: List[str] = None,
               additional_columns: Dict[str, Callable] = None) -> 'Table':

        if keep_columns is None:         # If no columns specified,
            keep_columns = self.columns  # return all columns

        if additional_columns is None:
            additional_columns = {}

        # New column names and types
        new_columns = keep_columns + list(additional_columns.keys())
        keep_types = [self.col2type(col) for col in keep_columns]

        # This is how to get the return type from a type annotation.
        # It will crash if `calculation` doesn't have a return type.
        add_types = [calculation.__annotations__['return']
                     for calculation in additional_columns.values()]

        # Create a new table for results
        new_table = Table(new_columns, keep_types + add_types)

        for row in self.rows:
            new_row = [row[column] for column in keep_columns]
            for column_name, calculation in additional_columns.items():
                new_row.append(calculation(row))
            new_table.insert(new_row)

        return new_table

    def where(self, predicate: WhereClause = lambda row: True) -> 'Table':
        """Return only the rows that satisfy the supplied predicate"""
        where_table = Table(self.columns, self.types)
        for row in self.rows:
            if predicate(row):
                values = [row[column] for column in self.columns]
                where_table.insert(values)
        return where_table

    def limit(self, num_rows: int) -> 'Table':
        """Return only the first `num_rows` rows"""
        limit_table = Table(self.columns, self.types)
        for i, row in enumerate(self.rows):
            if i >= num_rows:
                break
            values = [row[column] for column in self.columns]
            limit_table.insert(values)
        return limit_table

    def group_by(self,
                 group_by_columns: List[str],
                 aggregates: Dict[str, Callable],
                 having: HavingClause = lambda group: True) -> 'Table':

        grouped_rows = defaultdict(list)

        # Populate groups
        for row in self.rows:
            key = tuple(row[column] for column in group_by_columns)
            grouped_rows[key].append(row)

        # Result table consists of group_by columns and aggregates
        new_columns = group_by_columns + list(aggregates.keys())
        group_by_types = [self.col2type(col) for col in group_by_columns]
        aggregate_types = [agg.__annotations__['return']
                           for agg in aggregates.values()]
        result_table = Table(new_columns, group_by_types + aggregate_types)

        for key, rows in grouped_rows.items():
            if having(rows):
                new_row = list(key)
                for aggregate_name, aggregate_fn in aggregates.items():
                    new_row.append(aggregate_fn(rows))
                result_table.insert(new_row)

        return result_table

    def order_by(self, order: Callable[[Row], Any]) -> 'Table':
        new_table = self.select()       # make a copy
        new_table.rows.sort(key=order)
        return new_table

    def join(self, other_table: 'Table', left_join: bool = False) -> 'Table':

        join_on_columns = [c for c in self.columns           # columns in
                           if c in other_table.columns]      # both tables

        additional_columns = [c for c in other_table.columns # columns only
                              if c not in join_on_columns]   # in right table

        # all columns from left table + additional_columns from right table
        new_columns = self.columns + additional_columns
        new_types = self.types + [other_table.col2type(col)
                                  for col in additional_columns]

        join_table = Table(new_columns, new_types)

        for row in self.rows:
            def is_join(other_row):
                return all(other_row[c] == row[c] for c in join_on_columns)

            other_rows = other_table.where(is_join).rows

            # Each other row that matches this one produces a result row.
            for other_row in other_rows:
                join_table.insert([row[c] for c in self.columns] +
                                  [other_row[c] for c in additional_columns])

            # If no rows match and it's a left join, output with Nones.
            if left_join and not other_rows:
                join_table.insert([row[c] for c in self.columns] +
                                  [None for c in additional_columns])

        return join_table

def main():
    # Constructor requires column names and types
    users = Table(['user_id', 'name', 'num_friends'], [int, str, int])
    users.insert([0, "Hero", 0])
    users.insert([1, "Dunn", 2])
    users.insert([2, "Sue", 3])
    users.insert([3, "Chi", 3])
    users.insert([4, "Thor", 3])
    users.insert([5, "Clive", 2])
    users.insert([6, "Hicks", 3])
    users.insert([7, "Devin", 2])
    users.insert([8, "Kate", 2])
    users.insert([9, "Klein", 3])
    users.insert([10, "Jen", 1])
    
    assert len(users) == 11
    assert users[1]['name'] == 'Dunn'
    
    assert users[1]['num_friends'] == 2             # Original value
    
    users.update({'num_friends' : 3},               # Set num_friends = 3
                 lambda row: row['user_id'] == 1)   # in rows where user_id == 1
    
    assert users[1]['num_friends'] == 3             # Updated value
    
    # SELECT * FROM users;
    all_users = users.select()
    assert len(all_users) == 11
    
    # SELECT * FROM users LIMIT 2;
    two_users = users.limit(2)
    assert len(two_users) == 2
    
    # SELECT user_id FROM users;
    just_ids = users.select(keep_columns=["user_id"])
    assert just_ids.columns == ['user_id']
    
    # SELECT user_id FROM users WHERE name = 'Dunn';
    dunn_ids = (
        users
        .where(lambda row: row["name"] == "Dunn")
        .select(keep_columns=["user_id"])
    )
    assert len(dunn_ids) == 1
    assert dunn_ids[0] == {"user_id": 1}
    
    # SELECT LENGTH(name) AS name_length FROM users;
    def name_length(row) -> int: return len(row["name"])
    
    name_lengths = users.select(keep_columns=[],
                                additional_columns = {"name_length": name_length})
    assert name_lengths[0]['name_length'] == len("Hero")
    
    def min_user_id(rows) -> int:
        return min(row["user_id"] for row in rows)
    
    def length(rows) -> int:
        return len(rows)
    
    stats_by_length = (
        users
        .select(additional_columns={"name_length" : name_length})
        .group_by(group_by_columns=["name_length"],
                  aggregates={"min_user_id" : min_user_id,
                              "num_users" : length})
    )
    
    
    
    assert len(stats_by_length) == 3
    assert stats_by_length.columns == ["name_length", "min_user_id", "num_users"]
    
    def first_letter_of_name(row: Row) -> str:
        return row["name"][0] if row["name"] else ""
    
    def average_num_friends(rows: List[Row]) -> float:
        return sum(row["num_friends"] for row in rows) / len(rows)
    
    def enough_friends(rows: List[Row]) -> bool:
        return average_num_friends(rows) > 1
    
    avg_friends_by_letter = (
        users
        .select(additional_columns={'first_letter' : first_letter_of_name})
        .group_by(group_by_columns=['first_letter'],
                  aggregates={"avg_num_friends" : average_num_friends},
                  having=enough_friends)
    )
    
    
    assert len(avg_friends_by_letter) == 6
    assert {row['first_letter'] for row in avg_friends_by_letter} == \
           {"H", "D", "S", "C", "T", "K"}
    
    def sum_user_ids(rows: List[Row]) -> int:
        return sum(row["user_id"] for row in rows)
    
    user_id_sum = (
        users
        .where(lambda row: row["user_id"] > 1)
        .group_by(group_by_columns=[],
                  aggregates={ "user_id_sum" : sum_user_ids })
    )
    
    
    
    assert len(user_id_sum) == 1
    assert user_id_sum[0]["user_id_sum"] == 54
    
    friendliest_letters = (
        avg_friends_by_letter
        .order_by(lambda row: -row["avg_num_friends"])
        .limit(4)
    )
    
    
    assert len(friendliest_letters) == 4
    assert friendliest_letters[0]['first_letter'] in ['S', 'T']
    
    user_interests = Table(['user_id', 'interest'], [int, str])
    user_interests.insert([0, "SQL"])
    user_interests.insert([0, "NoSQL"])
    user_interests.insert([2, "SQL"])
    user_interests.insert([2, "MySQL"])
    
    sql_users = (
        users
        .join(user_interests)
        .where(lambda row: row["interest"] == "SQL")
        .select(keep_columns=["name"])
    )
    
    
    
    assert len(sql_users) == 2
    sql_user_names = {row["name"] for row in sql_users}
    assert sql_user_names == {"Hero", "Sue"}
    
    def count_interests(rows: List[Row]) -> int:
        """counts how many rows have non-None interests"""
        return len([row for row in rows if row["interest"] is not None])
    
    user_interest_counts = (
        users
        .join(user_interests, left_join=True)
        .group_by(group_by_columns=["user_id"],
                  aggregates={"num_interests" : count_interests })
    )
    
    likes_sql_user_ids = (
        user_interests
        .where(lambda row: row["interest"] == "SQL")
        .select(keep_columns=['user_id'])
    )
    
    likes_sql_user_ids.group_by(group_by_columns=[],
                                aggregates={ "min_user_id" : min_user_id })
    
    
    
    assert len(likes_sql_user_ids) == 2
    
    (
        user_interests
        .where(lambda row: row["interest"] == "SQL")
        .join(users)
        .select(["name"])
    )
    
    (
        user_interests
        .join(users)
        .where(lambda row: row["interest"] == "SQL")
        .select(["name"])
    )
    
if __name__ == "__main__": main()