DataTable

datatable.py

# (C) Tai Kedzierski
# LGPLv3

from collections import OrderedDict
from typing import Any, Union, List


class DataTable(list):
    def __init__(self, data:List[List[Any]]):
        """ Builds a table, after popping the first line which provides headers
        """
        list.__init__(data[1:])
        self.__headers = tuple(data[0])


    def column(self, name):
        """ Gets a column by name, and returns all values under it
        """
        return [R[name] for R in self]


    def headers(self):
        """ Return the headers of the table
        """
        return self.__headers

    
    def __getitem__(self, idx:Union[int,slice]) -> OrderedDict[str,Any]:
        """ Return data for the required row(s)
        
        If idx is an int, returns a row as a OrderedDict, first line of actual data being at index 0

        If idx is a slice, returns a DataTable
        """
        row = list.__getitem__(self, idx)

        if isinstance(idx, slice):
            subtable = list.__getitiem__(self, idx)
            subtable = [self.__headers, *subtable]
            return DataTable(subtable)

        query_row = TableRow()
        for i,k in enumerate(self.__headers):
            query_row[k] = row[i]
        return query_row


    def __iter__(self):
        """ Allow using column lookup per-row during iteration

        ```python
        spammy_meals = [meal for meal in data_table if "spam" in meal["ingredients"] ]
        ```
        """
        for row in list.__iter__(self):
            query_row = TableRow()
            for i,k in enumerate(self.__headers):
                query_row[k] = row[i]
            yield query_row


    def __contains__(self, __o: object) -> bool:
        """Will check if the item is in the DataTable itself but not the headers
        """
        for row in self:
            if __o in row:
                return True
        return False


class TableRow(OrderedDict):
    """ Wrapper for an ordered lookup dict

    Subscripting with a str will use a column name to lookup the value

    Subscripting with an int N will return the N-th item like for a list

    Subscripting with a slice will operate like for a list
    """
    def __init__(self):
        OrderedDict.__init__(self)


    def __getitem__(self, name):
        if isinstance(name, (int,slice)):
            return list(self)[name]

        return OrderedDict.__getitem__(self, name)