mirror of
https://github.com/shmilylty/OneForAll.git
synced 2026-08-26 04:47:48 +08:00
本地实现records和tablib
This commit is contained in:
JrD
+2
-2
@@ -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
|
||||
|
||||
|
||||
@@ -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 '<Record {}>'.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 '<RecordCollection size={} pending={}>'.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 '<Database open={}>'.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 '<Connection open={}>'.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)
|
||||
@@ -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
|
||||
@@ -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.<format> read or read/write properties
|
||||
|
||||
# Create Dataset.<format> read or read/write properties,
|
||||
# and Dataset.get_<format>/set_<format> 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()
|
||||
@@ -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 <newformats>`.
|
||||
|
||||
"""
|
||||
|
||||
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 '<dataset object>'
|
||||
|
||||
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"
|
||||
+1
-1
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user