# ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓ # ┃ ██████ ██████ ██████ █ █ █ █ █ █▄ ▀███ █ ┃ # ┃ ▄▄▄▄▄█ █▄▄▄▄▄ ▄▄▄▄▄█ ▀▀▀▀▀█▀▀▀▀▀ █ ▀▀▀▀▀█ ████████▌▐███ ███▄ ▀█ █ ▀▀▀▀▀ ┃ # ┃ █▀▀▀▀▀ █▀▀▀▀▀ █▀██▀▀ ▄▄▄▄▄ █ ▄▄▄▄▄█ ▄▄▄▄▄█ ████████▌▐███ █████▄ █ ▄▄▄▄▄ ┃ # ┃ █ ██████ █ ▀█▄ █ ██████ █ ███▌▐███ ███████▄ █ ┃ # ┣━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┫ # ┃ Copyright (c) 2017, the Perspective Authors. ┃ # ┃ ╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌ ┃ # ┃ This file is part of the Perspective library, distributed under the terms ┃ # ┃ of the [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0). ┃ # ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛ from datetime import datetime, date, timezone import numpy as np import pandas as pd import pyarrow as pa from pytest import fixture from random import random, randint, choice from faker import Faker # Our tests construct naive datetimes everywhere # so setting it here is an easy way to fix it globally. import os # Perspective used to support datetime.date and datetime.datetime # as Table() constructor arguments, but now we forward the parameters # directly to JSON.loads. So to make sure the tests dont need to be # so utterly transmogrified, we have this little hack :) import json os.environ["TZ"] = "UTC" def new_encoder(self, obj): if isinstance(obj, datetime): return str(obj) elif isinstance(obj, date): return str(obj) else: return old(self, obj) old = json.JSONEncoder.default json.JSONEncoder.default = new_encoder fake = Faker() def _make_date_time_index(size, time_unit): return pd.date_range("2000-01-01", periods=size, freq=time_unit) def _make_period_index(size, time_unit): return pd.period_range(start="2000", periods=size, freq=time_unit) def _make_dataframe(index, size=10): """Create a new random dataframe of `size` and with a DateTimeIndex of frequency `time_unit`. """ return pd.DataFrame( index=index, data={ "a": np.random.rand(size), "b": np.random.rand(size), "c": np.random.rand(size), "d": np.random.rand(size), }, ) class Util: @staticmethod def make_arrow(names, data, types=None, legacy=False): """Create an arrow binary that can be loaded and manipulated from memory. Args: names (list): a list of str column names data (list): a list of lists containing data for each column types (list): an optional list of `pyarrow.type` function references. Types will be inferred if not provided. legacy (bool): if True, use legacy IPC format (pre-pyarrow 0.15). Defaults to False. Returns: bytes : a bytes object containing the arrow-serialized output. """ stream = pa.BufferOutputStream() arrays = [] for idx, column in enumerate(data): # only apply types if array is present kwargs = {} if types: kwargs["type"] = types[idx] arrays.append(pa.array(column, **kwargs)) batch = pa.RecordBatch.from_arrays(arrays, names) table = pa.Table.from_batches([batch]) writer = pa.RecordBatchStreamWriter( stream, table.schema ) writer.write_table(table) writer.close() return stream.getvalue().to_pybytes() @staticmethod def make_arrow_from_pandas(df, schema=None, legacy=False): """Create a pyarrow Table from a Pandas dataframe. Args: df (:obj:`pandas.DataFrame`) schema (:obj:`pyarrow.Schema`) legacy (bool): unused; retained for backwards compatibility. Returns: pyarrow.Table """ return pa.Table.from_pandas(df, schema=schema) @staticmethod def make_dictionary_arrow(names, data, types=None, legacy=False): """Create an arrow binary that can be loaded and manipulated from memory, with each column being a dictionary array of `str` values and `int` indices. Args: names (list): a list of str column names data (list:tuple): a list of tuples, the first value being a list of indices, and the second value being a list of values. types (list:list:pyarrow.func): a list of lists, containing the indices type and dictionary value type for each array. legacy (bool): if True, use legacy IPC format (pre-pyarrow 0.15). Defaults to False. Returns: bytes : a bytes object containing the arrow-serialized output. """ stream = pa.BufferOutputStream() arrays = [] for idx, column in enumerate(data): indice_type = pa.int64() value_type = pa.string() if types is not None: indice_type = types[idx][0] value_type = types[idx][1] indices = pa.array(column[0], type=indice_type) values = pa.array(column[1], type=value_type) parray = pa.DictionaryArray.from_arrays(indices, values) arrays.append(parray) batch = pa.RecordBatch.from_arrays(arrays, names) table = pa.Table.from_batches([batch]) writer = pa.RecordBatchStreamWriter( stream, table.schema ) writer.write_table(table) writer.close() return stream.getvalue().to_pybytes() @staticmethod def to_timestamp(obj): """Return an integer timestamp based on a date/datetime object. A `date` is a timezone-agnostic calendar day and serializes to epoch ms at *UTC* midnight regardless of the host process timezone, so it is converted here in UTC. A naive `datetime` is a wall-clock reading in the process-local timezone (callers express expected instants this way, e.g. `aware.astimezone(TZ).replace(tzinfo=None)`), so it is converted via local `timestamp()`. """ classname = obj.__class__.__name__ if classname == "date": return int( datetime( obj.year, obj.month, obj.day, tzinfo=timezone.utc ).timestamp() * 1000 ) elif classname == "datetime": return int(obj.timestamp() * 1000) else: return -1 @staticmethod def make_dataframe(size=10, freq="D"): index = _make_date_time_index(size, freq) return _make_dataframe(index, size) @staticmethod def make_period_dataframe(size=10): index = _make_period_index(size, "M") return _make_dataframe(index, size) @staticmethod def make_series(size=10, freq="D"): index = _make_date_time_index(size, freq) return pd.Series(data=np.random.rand(size), index=index) class Sentinel(object): """Generic sentinel class for testing side-effectful code in Python 2 and 3. """ def __init__(self, value): self.value = value def get(self): return self.value def set(self, new_value): self.value = new_value @fixture() def sentinel(): """Pass `sentinel` into a test and call it with `value` to create a new instance of the Sentinel class. Example: >>> def test_with_sentinel(self, sentinel): >>> s = sentinel(True) >>> s.set(False) >>> s.get() # returns False """ def _sentinel(value): return Sentinel(value) return _sentinel @fixture def util(): """Pass the `Util` class in to a test.""" return Util @fixture def superstore(count=100): data = [] for id in range(count): dat = {} dat["Row ID"] = id dat["Order ID"] = "{}-{}".format(fake.ein(), fake.zipcode()) dat["Order Date"] = fake.date_this_year() dat["Ship Date"] = fake.date_between_dates(dat["Order Date"]).strftime( "%Y-%m-%d" ) dat["Order Date"] = dat["Order Date"].strftime("%Y-%m-%d") dat["Ship Mode"] = choice(["First Class", "Standard Class", "Second Class"]) dat["Ship Mode"] = choice(["First Class", "Standard Class", "Second Class"]) dat["Customer ID"] = fake.zipcode() dat["Segment"] = choice(["A", "B", "C", "D"]) dat["Country"] = "US" dat["City"] = fake.city() dat["State"] = fake.state() dat["Postal Code"] = fake.zipcode() dat["Region"] = choice(["Region %d" % i for i in range(5)]) dat["Product ID"] = fake.bban() sector = choice(["Industrials", "Technology", "Financials"]) industry = choice(["A", "B", "C"]) dat["Category"] = sector dat["Sub-Category"] = industry dat["Sales"] = randint(1, 100) * 100 dat["Quantity"] = randint(1, 100) * 10 dat["Discount"] = round(random() * 100, 2) dat["Profit"] = round(random() * 1000, 2) data.append(dat) return pd.DataFrame(data)