diff --git a/common/database.py b/common/database.py index 5fde260..415ca0b 100644 --- a/common/database.py +++ b/common/database.py @@ -2,9 +2,9 @@ SQLite database initialization and operation """ -import records +from common import records -from records import Connection +from common.records import Connection from config.log import logger from config import settings diff --git a/common/records.py b/common/records.py new file mode 100644 index 0000000..867a997 --- /dev/null +++ b/common/records.py @@ -0,0 +1,440 @@ +# -*- coding: utf-8 -*- + +import os +from sys import stdout +from collections import OrderedDict +from contextlib import contextmanager +from inspect import isclass + +from .tablib import tablib +from sqlalchemy import create_engine, exc, inspect, text + +DATABASE_URL = os.environ.get('DATABASE_URL') + + +def isexception(obj): + """Given an object, return a boolean indicating whether it is an instance + or subclass of :py:class:`Exception`. + """ + if isinstance(obj, Exception): + return True + if isclass(obj) and issubclass(obj, Exception): + return True + return False + + +class Record(object): + """A row, from a query, from a database.""" + __slots__ = ('_keys', '_values') + + def __init__(self, keys, values): + self._keys = keys + self._values = values + + # Ensure that lengths match properly. + assert len(self._keys) == len(self._values) + + def keys(self): + """Returns the list of column names from the query.""" + return self._keys + + def values(self): + """Returns the list of values from the query.""" + return self._values + + def __repr__(self): + return ''.format(self.export('json')[1:-1]) + + def __getitem__(self, key): + # Support for index-based lookup. + if isinstance(key, int): + return self.values()[key] + + # Support for string-based lookup. + if key in self.keys(): + i = self.keys().index(key) + if self.keys().count(key) > 1: + raise KeyError("Record contains multiple '{}' fields.".format(key)) + return self.values()[i] + + raise KeyError("Record contains no '{}' field.".format(key)) + + def __getattr__(self, key): + try: + return self[key] + except KeyError as e: + raise AttributeError(e) + + def __dir__(self): + standard = dir(super(Record, self)) + # Merge standard attrs with generated ones (from column names). + return sorted(standard + [str(k) for k in self.keys()]) + + def get(self, key, default=None): + """Returns the value for a given key, or default.""" + try: + return self[key] + except KeyError: + return default + + def as_dict(self, ordered=False): + """Returns the row as a dictionary, as ordered.""" + items = zip(self.keys(), self.values()) + + return OrderedDict(items) if ordered else dict(items) + + @property + def dataset(self): + """A Tablib Dataset containing the row.""" + data = tablib.Dataset() + data.headers = self.keys() + + row = _reduce_datetimes(self.values()) + data.append(row) + + return data + + def export(self, format, **kwargs): + """Exports the row to the given format.""" + return self.dataset.export(format, **kwargs) + + +class RecordCollection(object): + """A set of excellent Records from a query.""" + + def __init__(self, rows): + self._rows = rows + self._all_rows = [] + self.pending = True + + def __repr__(self): + return ''.format(len(self), self.pending) + + def __iter__(self): + """Iterate over all rows, consuming the underlying generator + only when necessary.""" + i = 0 + while True: + # Other code may have iterated between yields, + # so always check the cache. + if i < len(self): + yield self[i] + else: + # Throws StopIteration when done. + # Prevent StopIteration bubbling from generator, following https://www.python.org/dev/peps/pep-0479/ + try: + yield next(self) + except StopIteration: + return + i += 1 + + def next(self): + return self.__next__() + + def __next__(self): + try: + nextrow = next(self._rows) + self._all_rows.append(nextrow) + return nextrow + except StopIteration: + self.pending = False + raise StopIteration('RecordCollection contains no more rows.') + + def __getitem__(self, key): + is_int = isinstance(key, int) + + # Convert RecordCollection[1] into slice. + if is_int: + key = slice(key, key + 1) + + while len(self) < key.stop or key.stop is None: + try: + next(self) + except StopIteration: + break + + rows = self._all_rows[key] + if is_int: + return rows[0] + else: + return RecordCollection(iter(rows)) + + def __len__(self): + return len(self._all_rows) + + def export(self, format, **kwargs): + """Export the RecordCollection to a given format (courtesy of Tablib).""" + return self.dataset.export(format, **kwargs) + + @property + def dataset(self): + """A Tablib Dataset representation of the RecordCollection.""" + # Create a new Tablib Dataset. + data = tablib.Dataset() + + # If the RecordCollection is empty, just return the empty set + # Check number of rows by typecasting to list + if len(list(self)) == 0: + return data + + # Set the column names as headers on Tablib Dataset. + first = self[0] + + data.headers = first.keys() + for row in self.all(): + row = _reduce_datetimes(row.values()) + data.append(row) + + return data + + def all(self, as_dict=False, as_ordereddict=False): + """Returns a list of all rows for the RecordCollection. If they haven't + been fetched yet, consume the iterator and cache the results.""" + + # By calling list it calls the __iter__ method + rows = list(self) + + if as_dict: + return [r.as_dict() for r in rows] + elif as_ordereddict: + return [r.as_dict(ordered=True) for r in rows] + + return rows + + def as_dict(self, ordered=False): + return self.all(as_dict=not (ordered), as_ordereddict=ordered) + + def first(self, default=None, as_dict=False, as_ordereddict=False): + """Returns a single record for the RecordCollection, or `default`. If + `default` is an instance or subclass of Exception, then raise it + instead of returning it.""" + + # Try to get a record, or return/raise default. + try: + record = self[0] + except IndexError: + if isexception(default): + raise default + return default + + # Cast and return. + if as_dict: + return record.as_dict() + elif as_ordereddict: + return record.as_dict(ordered=True) + else: + return record + + def one(self, default=None, as_dict=False, as_ordereddict=False): + """Returns a single record for the RecordCollection, ensuring that it + is the only record, or returns `default`. If `default` is an instance + or subclass of Exception, then raise it instead of returning it.""" + + # Ensure that we don't have more than one row. + try: + self[1] + except IndexError: + return self.first(default=default, as_dict=as_dict, as_ordereddict=as_ordereddict) + else: + raise ValueError('RecordCollection contained more than one row. ' + 'Expects only one row when using ' + 'RecordCollection.one') + + def scalar(self, default=None): + """Returns the first column of the first row, or `default`.""" + row = self.one() + return row[0] if row else default + + +class Database(object): + """A Database. Encapsulates a url and an SQLAlchemy engine with a pool of + connections. + """ + + def __init__(self, db_url=None, **kwargs): + # If no db_url was provided, fallback to $DATABASE_URL. + self.db_url = db_url or DATABASE_URL + + if not self.db_url: + raise ValueError('You must provide a db_url.') + + # Create an engine. + self._engine = create_engine(self.db_url, **kwargs) + self.open = True + + def close(self): + """Closes the Database.""" + self._engine.dispose() + self.open = False + + def __enter__(self): + return self + + def __exit__(self, exc, val, traceback): + self.close() + + def __repr__(self): + return ''.format(self.open) + + def get_table_names(self, internal=False): + """Returns a list of table names for the connected database.""" + + # Setup SQLAlchemy for Database inspection. + return inspect(self._engine).get_table_names() + + def get_connection(self): + """Get a connection to this Database. Connections are retrieved from a + pool. + """ + if not self.open: + raise exc.ResourceClosedError('Database closed.') + + return Connection(self._engine.connect()) + + def query(self, query, fetchall=False, **params): + """Executes the given SQL query against the Database. Parameters can, + optionally, be provided. Returns a RecordCollection, which can be + iterated over to get result rows as dictionaries. + """ + with self.get_connection() as conn: + return conn.query(query, fetchall, **params) + + def bulk_query(self, query, *multiparams): + """Bulk insert or update.""" + + with self.get_connection() as conn: + conn.bulk_query(query, *multiparams) + + def query_file(self, path, fetchall=False, **params): + """Like Database.query, but takes a filename to load a query from.""" + + with self.get_connection() as conn: + return conn.query_file(path, fetchall, **params) + + def bulk_query_file(self, path, *multiparams): + """Like Database.bulk_query, but takes a filename to load a query from.""" + + with self.get_connection() as conn: + conn.bulk_query_file(path, *multiparams) + + @contextmanager + def transaction(self): + """A context manager for executing a transaction on this Database.""" + + conn = self.get_connection() + tx = conn.transaction() + try: + yield conn + tx.commit() + except: + tx.rollback() + finally: + conn.close() + + +class Connection(object): + """A Database connection.""" + + def __init__(self, connection): + self._conn = connection + self.open = not connection.closed + + def close(self): + self._conn.close() + self.open = False + + def __enter__(self): + return self + + def __exit__(self, exc, val, traceback): + self.close() + + def __repr__(self): + return ''.format(self.open) + + def query(self, query, fetchall=False, **params): + """Executes the given SQL query against the connected Database. + Parameters can, optionally, be provided. Returns a RecordCollection, + which can be iterated over to get result rows as dictionaries. + """ + + # Execute the given query. + cursor = self._conn.execute(text(query), **params) # TODO: PARAMS GO HERE + + # Row-by-row Record generator. + row_gen = (Record(cursor.keys(), row) for row in cursor) + + # Convert psycopg2 results to RecordCollection. + results = RecordCollection(row_gen) + + # Fetch all results if desired. + if fetchall: + results.all() + + return results + + def bulk_query(self, query, *multiparams): + """Bulk insert or update.""" + + self._conn.execute(text(query), *multiparams) + + def query_file(self, path, fetchall=False, **params): + """Like Connection.query, but takes a filename to load a query from.""" + + # If path doesn't exists + if not os.path.exists(path): + raise IOError("File '{}' not found!".format(path)) + + # If it's a directory + if os.path.isdir(path): + raise IOError("'{}' is a directory!".format(path)) + + # Read the given .sql file into memory. + with open(path) as f: + query = f.read() + + # Defer processing to self.query method. + return self.query(query=query, fetchall=fetchall, **params) + + def bulk_query_file(self, path, *multiparams): + """Like Connection.bulk_query, but takes a filename to load a query + from. + """ + + # If path doesn't exists + if not os.path.exists(path): + raise IOError("File '{}'' not found!".format(path)) + + # If it's a directory + if os.path.isdir(path): + raise IOError("'{}' is a directory!".format(path)) + + # Read the given .sql file into memory. + with open(path) as f: + query = f.read() + + self._conn.execute(text(query), *multiparams) + + def transaction(self): + """Returns a transaction object. Call ``commit`` or ``rollback`` + on the returned object as appropriate.""" + + return self._conn.begin() + + +def _reduce_datetimes(row): + """Receives a row, converts datetimes to strings.""" + + row = list(row) + + for i in range(len(row)): + if hasattr(row[i], 'isoformat'): + row[i] = row[i].isoformat() + return tuple(row) + + +def print_bytes(content): + try: + stdout.buffer.write(content) + except AttributeError: + stdout.write(content) diff --git a/common/tablib/__init__.py b/common/tablib/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/common/tablib/format.py b/common/tablib/format.py new file mode 100644 index 0000000..b52ceb9 --- /dev/null +++ b/common/tablib/format.py @@ -0,0 +1,120 @@ +""" Tablib - JSON Support +""" +import decimal +import json +from uuid import UUID + +from . import tablib + + +def serialize_objects_handler(obj): + if isinstance(obj, (decimal.Decimal, UUID)): + return str(obj) + elif hasattr(obj, 'isoformat'): + return obj.isoformat() + else: + return obj + + +class JSONFormat: + title = 'json' + extensions = ('json',) + + @classmethod + def export_set(cls, dataset): + """Returns JSON representation of Dataset.""" + return json.dumps(dataset.dict, default=serialize_objects_handler) + + @classmethod + def export_book(cls, databook): + """Returns JSON representation of Databook.""" + return json.dumps(databook._package(), default=serialize_objects_handler) + + @classmethod + def import_set(cls, dset, in_stream): + """Returns dataset from JSON stream.""" + + dset.wipe() + dset.dict = json.load(in_stream) + + @classmethod + def import_book(cls, dbook, in_stream): + """Returns databook from JSON stream.""" + + dbook.wipe() + for sheet in json.load(in_stream): + data = tablib.Dataset() + data.title = sheet['title'] + data.dict = sheet['data'] + dbook.add_sheet(data) + + @classmethod + def detect(cls, stream): + """Returns True if given stream is valid JSON.""" + try: + json.load(stream) + return True + except (TypeError, ValueError): + return False + + +""" Tablib - *SV Support. +""" + +import csv +from io import StringIO + + +class CSVFormat: + title = 'csv' + extensions = ('csv',) + + DEFAULT_DELIMITER = ',' + + @classmethod + def export_stream_set(cls, dataset, **kwargs): + """Returns CSV representation of Dataset as file-like.""" + stream = StringIO() + + kwargs.setdefault('delimiter', cls.DEFAULT_DELIMITER) + + _csv = csv.writer(stream, **kwargs) + + for row in dataset._package(dicts=False): + _csv.writerow(row) + + stream.seek(0) + return stream + + @classmethod + def export_set(cls, dataset, **kwargs): + """Returns CSV representation of Dataset.""" + stream = cls.export_stream_set(dataset, **kwargs) + return stream.getvalue() + + @classmethod + def import_set(cls, dset, in_stream, headers=True, **kwargs): + """Returns dataset from CSV stream.""" + + dset.wipe() + + kwargs.setdefault('delimiter', cls.DEFAULT_DELIMITER) + + rows = csv.reader(in_stream, **kwargs) + for i, row in enumerate(rows): + + if (i == 0) and (headers): + dset.headers = row + elif row: + if i > 0 and len(row) < dset.width: + row += [''] * (dset.width - len(row)) + dset.append(row) + + @classmethod + def detect(cls, stream, delimiter=None): + """Returns True if given stream is valid CSV.""" + try: + csv.Sniffer().sniff(stream.read(1024), delimiters=delimiter or cls.DEFAULT_DELIMITER) + return True + except Exception: + return False diff --git a/common/tablib/registry.py b/common/tablib/registry.py new file mode 100644 index 0000000..3d908dc --- /dev/null +++ b/common/tablib/registry.py @@ -0,0 +1,28 @@ +""" Tablib - formats +""" +from collections import OrderedDict +from .format import JSONFormat, CSVFormat + + +class Registry: + _formats = OrderedDict() + + def register(self, key, format_or_path): + # Create Databook. read or read/write properties + + # Create Dataset. read or read/write properties, + # and Dataset.get_/set_ methods. + self._formats[key] = format_or_path + + def register_builtins(self): + # Registration ordering matters for autodetection. + self.register('csv', CSVFormat()) + self.register('json', JSONFormat()) + + def get_format(self, key): + if key not in self._formats: + raise Exception("OneForAll has no format '%s'." % key) + return self._formats[key] + + +registry = Registry() diff --git a/common/tablib/tablib.py b/common/tablib/tablib.py new file mode 100644 index 0000000..dd53137 --- /dev/null +++ b/common/tablib/tablib.py @@ -0,0 +1,530 @@ +from collections import OrderedDict +from io import BytesIO, StringIO +from .registry import registry + +def normalize_input(stream): + """ + Accept either a str/bytes stream or a file-like object and always return a + file-like object. + """ + if isinstance(stream, str): + return StringIO(stream) + elif isinstance(stream, bytes): + return BytesIO(stream) + return stream + +def import_set(stream, format=None, **kwargs): + """Return dataset of given stream (file-like object, string, or bytestring).""" + + return Dataset().load(normalize_input(stream), format, **kwargs) + + +def detect_format(stream): + """Return format name of given stream (file-like object, string, or bytestring).""" + stream = normalize_input(stream) + fmt_title = None + for fmt in registry.formats(): + try: + if fmt.detect(stream): + fmt_title = fmt.title + break + except AttributeError: + pass + finally: + if hasattr(stream, 'seek'): + stream.seek(0) + return fmt_title + + +def get_format(format): + """ + Determine if the format is available + :param format: + :return: + """ + + +class Row: + """Internal Row object. Mainly used for filtering.""" + + __slots__ = ['_row', 'tags'] + + def __init__(self, row=None, tags=None): + if tags is None: + tags = list() + if row is None: + row = list() + self._row = list(row) + self.tags = list(tags) + + def __iter__(self): + return (col for col in self._row) + + def __len__(self): + return len(self._row) + + def __repr__(self): + return repr(self._row) + + def __getitem__(self, i): + return self._row[i] + + def __setitem__(self, i, value): + self._row[i] = value + + def __delitem__(self, i): + del self._row[i] + + def __getstate__(self): + + slots = dict() + + for slot in self.__slots__: + attribute = getattr(self, slot) + slots[slot] = attribute + + return slots + + def __setstate__(self, state): + for (k, v) in list(state.items()): + setattr(self, k, v) + + def rpush(self, value): + self.insert(len(self._row), value) + + def append(self, value): + self.rpush(value) + + def insert(self, index, value): + self._row.insert(index, value) + + def __contains__(self, item): + return (item in self._row) + + @property + def tuple(self): + """Tuple representation of :class:`Row`.""" + return tuple(self._row) + + +class Dataset: + """The :class:`Dataset` object is the heart of Tablib. It provides all core + functionality. + + Usually you create a :class:`Dataset` instance in your main module, and append + rows as you collect data. :: + + data = tablib.Dataset() + data.headers = ('name', 'age') + + for (name, age) in some_collector(): + data.append((name, age)) + + + Setting columns is similar. The column data length must equal the + current height of the data and headers must be set. :: + + data = tablib.Dataset() + data.headers = ('first_name', 'last_name') + + data.append(('John', 'Adams')) + data.append(('George', 'Washington')) + + data.append_col((90, 67), header='age') + + + You can also set rows and headers upon instantiation. This is useful if + dealing with dozens or hundreds of :class:`Dataset` objects. :: + + headers = ('first_name', 'last_name') + data = [('John', 'Adams'), ('George', 'Washington')] + + data = tablib.Dataset(*data, headers=headers) + + :param \\*args: (optional) list of rows to populate Dataset + :param headers: (optional) list strings for Dataset header row + :param title: (optional) string to use as title of the Dataset + + + .. admonition:: Format Attributes Definition + + If you look at the code, the various output/import formats are not + defined within the :class:`Dataset` object. To add support for a new format, see + :ref:`Adding New Formats `. + + """ + + def __init__(self, *args, **kwargs): + self._data = list(Row(arg) for arg in args) + self.__headers = None + + # ('title', index) tuples + self._separators = [] + + # (column, callback) tuples + self._formatters = [] + + self.headers = kwargs.get('headers') + + self.title = kwargs.get('title') + + def __len__(self): + return self.height + + def __getitem__(self, key): + if isinstance(key, str): + if key in self.headers: + pos = self.headers.index(key) # get 'key' index from each data + return [row[pos] for row in self._data] + else: + raise KeyError + else: + _results = self._data[key] + if isinstance(_results, Row): + return _results.tuple + else: + return [result.tuple for result in _results] + + def __setitem__(self, key, value): + self._validate(value) + self._data[key] = Row(value) + + def __delitem__(self, key): + if isinstance(key, str): + + if key in self.headers: + + pos = self.headers.index(key) + del self.headers[pos] + + for i, row in enumerate(self._data): + del row[pos] + self._data[i] = row + else: + raise KeyError + else: + del self._data[key] + + def __repr__(self): + try: + return '<%s dataset>' % (self.title.lower()) + except AttributeError: + return '' + + def __str__(self): + result = [] + + # Add str representation of headers. + if self.__headers: + result.append([str(h) for h in self.__headers]) + + # Add str representation of rows. + result.extend(list(map(str, row)) for row in self._data) + + lens = [list(map(len, row)) for row in result] + field_lens = list(map(max, zip(*lens))) + + # delimiter between header and data + if self.__headers: + result.insert(1, ['-' * length for length in field_lens]) + + format_string = '|'.join('{%s:%s}' % item for item in enumerate(field_lens)) + + return '\n'.join(format_string.format(*row) for row in result) + + # --------- + # Internals + # --------- + + def _get_in_format(self, fmt_key, **kwargs): + return registry.get_format(fmt_key).export_set(self, **kwargs) + + def _set_in_format(self, fmt_key, in_stream, **kwargs): + in_stream = normalize_input(in_stream) + return registry.get_format(fmt_key).import_set(self, in_stream, **kwargs) + + def _validate(self, row=None, col=None, safety=False): + """Assures size of every row in dataset is of proper proportions.""" + if row: + is_valid = (len(row) == self.width) if self.width else True + elif col: + if len(col) < 1: + is_valid = True + else: + is_valid = (len(col) == self.height) if self.height else True + else: + is_valid = all(len(x) == self.width for x in self._data) + + if is_valid: + return True + else: + if not safety: + raise InvalidDimensions + return False + + def _package(self, dicts=True, ordered=True): + """Packages Dataset into lists of dictionaries for transmission.""" + # TODO: Dicts default to false? + + _data = list(self._data) + + if ordered: + dict_pack = OrderedDict + else: + dict_pack = dict + + # Execute formatters + if self._formatters: + for row_i, row in enumerate(_data): + for col, callback in self._formatters: + try: + if col is None: + for j, c in enumerate(row): + _data[row_i][j] = callback(c) + else: + _data[row_i][col] = callback(row[col]) + except IndexError: + raise InvalidDatasetIndex + + if self.headers: + if dicts: + data = [dict_pack(list(zip(self.headers, data_row))) for data_row in _data] + else: + data = [list(self.headers)] + list(_data) + else: + data = [list(row) for row in _data] + + return data + + def _get_headers(self): + """An *optional* list of strings to be used for header rows and attribute names. + + This must be set manually. The given list length must equal :class:`Dataset.width`. + + """ + return self.__headers + + def _set_headers(self, collection): + """Validating headers setter.""" + self._validate(collection) + if collection: + try: + self.__headers = list(collection) + except TypeError: + raise TypeError + else: + self.__headers = None + + headers = property(_get_headers, _set_headers) + + def _get_dict(self): + """A native Python representation of the :class:`Dataset` object. If headers have + been set, a list of Python dictionaries will be returned. If no headers have been set, + a list of tuples (rows) will be returned instead. + + A dataset object can also be imported by setting the `Dataset.dict` attribute: :: + + data = tablib.Dataset() + data.dict = [{'age': 90, 'first_name': 'Kenneth', 'last_name': 'Reitz'}] + + """ + return self._package() + + def _set_dict(self, pickle): + """A native Python representation of the Dataset object. If headers have been + set, a list of Python dictionaries will be returned. If no headers have been + set, a list of tuples (rows) will be returned instead. + + A dataset object can also be imported by setting the :class:`Dataset.dict` attribute. :: + + data = tablib.Dataset() + data.dict = [{'age': 90, 'first_name': 'Kenneth', 'last_name': 'Reitz'}] + + """ + + if not len(pickle): + return + + # if list of rows + if isinstance(pickle[0], list): + self.wipe() + for row in pickle: + self.append(Row(row)) + + # if list of objects + elif isinstance(pickle[0], dict): + self.wipe() + self.headers = list(pickle[0].keys()) + for row in pickle: + self.append(Row(list(row.values()))) + else: + raise UnsupportedFormat + + dict = property(_get_dict, _set_dict) + + def _clean_col(self, col): + """Prepares the given column for insert/append.""" + + col = list(col) + + if self.headers: + header = [col.pop(0)] + else: + header = [] + + if len(col) == 1 and hasattr(col[0], '__call__'): + col = list(map(col[0], self._data)) + col = tuple(header + col) + + return col + + @property + def height(self): + """The number of rows currently in the :class:`Dataset`. + Cannot be directly modified. + """ + return len(self._data) + + @property + def width(self): + """The number of columns currently in the :class:`Dataset`. + Cannot be directly modified. + """ + + try: + return len(self._data[0]) + except IndexError: + try: + return len(self.headers) + except TypeError: + return 0 + + def export(self, format, **kwargs): + """ + Export :class:`Dataset` object to `format`. + + :param \\*\\*kwargs: (optional) custom configuration to the format `export_set`. + """ + fmt = registry.get_format(format) + if not hasattr(fmt, 'export_set'): + raise Exception('Format {} cannot be exported.'.format(format)) + + return fmt.export_set(self, **kwargs) + + # ---- + # Rows + # ---- + + def insert(self, index, row, tags=None): + """Inserts a row to the :class:`Dataset` at the given index. + + Rows inserted must be the correct size (height or width). + + The default behaviour is to insert the given row to the :class:`Dataset` + object at the given index. + """ + + if tags is None: + tags = list() + self._validate(row) + self._data.insert(index, Row(row, tags=tags)) + + def rpush(self, row, tags=None): + """Adds a row to the end of the :class:`Dataset`. + See :class:`Dataset.insert` for additional documentation. + """ + + if tags is None: + tags = list() + self.insert(self.height, row=row, tags=tags) + + def append(self, row, tags=None): + """Adds a row to the :class:`Dataset`. + See :class:`Dataset.insert` for additional documentation. + """ + + if tags is None: + tags = list() + self.rpush(row, tags) + + def extend(self, rows, tags=None): + """Adds a list of rows to the :class:`Dataset` using + :class:`Dataset.append` + """ + + if tags is None: + tags = list() + for row in rows: + self.append(row, tags) + + # ---- + # Misc + # ---- + + def add_formatter(self, col, handler): + """Adds a formatter to the :class:`Dataset`. + + .. versionadded:: 0.9.5 + + :param col: column to. Accepts index int or header str. + :param handler: reference to callback function to execute against + each cell value. + """ + + if isinstance(col, str): + if col in self.headers: + col = self.headers.index(col) # get 'key' index from each data + else: + raise KeyError + + if not col > self.width: + self._formatters.append((col, handler)) + else: + raise InvalidDatasetIndex + + return True + + def remove_duplicates(self): + """Removes all duplicate rows from the :class:`Dataset` object + while maintaining the original order.""" + seen = set() + self._data[:] = [row for row in self._data if not (tuple(row) in seen or seen.add(tuple(row)))] + + def wipe(self): + """Removes all content and headers from the :class:`Dataset` object.""" + self._data = list() + self.__headers = None + + def load(self, in_stream, format, **kwargs): + """ + Import `in_stream` to the :class:`Databook` object using the `format`. + `in_stream` can be a file-like object, a string, or a bytestring. + + :param \\*\\*kwargs: (optional) custom configuration to the format `import_book`. + """ + + stream = normalize_input(in_stream) + if not format: + format = detect_format(stream) + + fmt = registry.get_format(format) + if not hasattr(fmt, 'import_book'): + raise UnsupportedFormat('Format {} cannot be loaded.'.format(format)) + + fmt.import_book(self, stream, **kwargs) + return self + + +registry.register_builtins() + + +class InvalidDimensions(Exception): + "Invalid size" + + +class InvalidDatasetIndex(Exception): + "Outside of Dataset size" + + +class UnsupportedFormat(NotImplementedError): + "Format is not supported" diff --git a/common/utils.py b/common/utils.py index 0ca8474..fd70b8b 100644 --- a/common/utils.py +++ b/common/utils.py @@ -13,7 +13,7 @@ from stat import S_IXUSR import tenacity import requests from pathlib import Path -from records import Record, RecordCollection +from common.records import Record, RecordCollection from dns.resolver import Resolver from common.domain import Domain @@ -731,5 +731,5 @@ def decode_resp_text(resp): content = str(content, encoding='gb18030', errors='strict') except (LookupError, TypeError, UnicodeError): # 最后尝试自动解码 - content = str(content, errors='replace') + content = str(content, errors='replace') return content diff --git a/takeover.py b/takeover.py index 8239034..721b2c4 100644 --- a/takeover.py +++ b/takeover.py @@ -13,7 +13,7 @@ from threading import Thread from queue import Queue import fire -from tablib import Dataset +from common.tablib.tablib import Dataset from tqdm import tqdm from config.log import logger