Support for numpy types, closes #11 (#12)

This commit is contained in:
Simon Willison 2019-02-23 20:02:19 -08:00 committed by GitHub
commit e615d22c55
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 102 additions and 0 deletions

View file

@ -12,6 +12,8 @@ script:
- sudo apt-get update && sudo apt-get install sqlite3
- pip install -U pip wheel
- pip install .[test]
# Only run the numpy/pandas tests on Python 3.7:
- python -c "import sys; print(sys.version_info.minor == 7)" | grep True > /dev/null && pip install pandas || true
- pytest
cache:

View file

@ -5,6 +5,11 @@ import itertools
import json
import pathlib
try:
import numpy as np
except ImportError:
np = None
Column = namedtuple(
"Column", ("cid", "name", "type", "notnull", "default_value", "is_pk")
)
@ -70,6 +75,23 @@ class Database:
datetime.time: "TEXT",
None.__class__: "TEXT",
}
# If numpy is available, add more types
if np:
col_type_mapping.update(
{
np.int8: "INTEGER",
np.int16: "INTEGER",
np.int32: "INTEGER",
np.int64: "INTEGER",
np.uint8: "INTEGER",
np.uint16: "INTEGER",
np.uint32: "INTEGER",
np.uint64: "INTEGER",
np.float16: "FLOAT",
np.float32: "FLOAT",
np.float64: "FLOAT",
}
)
columns_sql = ",\n".join(
" [{col_name}] {col_type}{primary_key}{references}".format(
col_name=col_name,

View file

@ -5,6 +5,11 @@ import pathlib
import pytest
import json
try:
import pandas as pd
except ImportError:
pd = None
def test_create_table(fresh_db):
assert [] == fresh_db.table_names()
@ -218,3 +223,76 @@ def test_works_with_pathlib_path(tmpdir):
db = Database(path)
db["demo"].insert_all([{"foo": 1}])
assert 1 == db["demo"].count
@pytest.mark.skipif(pd is None, reason="pandas and numpy are not installed")
def test_create_table_numpy(fresh_db):
import numpy as np
df = pd.DataFrame({"col 1": range(3), "col 2": range(3)})
fresh_db["pandas"].insert_all(df.to_dict(orient="records"))
assert [
{"col 1": 0, "col 2": 0},
{"col 1": 1, "col 2": 1},
{"col 1": 2, "col 2": 2},
] == list(fresh_db["pandas"].rows)
# Now try all the different types
df = pd.DataFrame(
{
"np.int8": [-8],
"np.int16": [-16],
"np.int32": [-32],
"np.int64": [-64],
"np.uint8": [8],
"np.uint16": [16],
"np.uint32": [32],
"np.uint64": [64],
"np.float16": [16.5],
"np.float32": [32.5],
"np.float64": [64.5],
}
)
df = df.astype(
{
"np.int8": "int8",
"np.int16": "int16",
"np.int32": "int32",
"np.int64": "int64",
"np.uint8": "uint8",
"np.uint16": "uint16",
"np.uint32": "uint32",
"np.uint64": "uint64",
"np.float16": "float16",
"np.float32": "float32",
"np.float64": "float64",
}
)
assert [
"int8",
"int16",
"int32",
"int64",
"uint8",
"uint16",
"uint32",
"uint64",
"float16",
"float32",
"float64",
] == [str(t) for t in df.dtypes]
fresh_db["types"].insert_all(df.to_dict(orient="records"))
assert [
{
"np.float16": 16.5,
"np.float32": 32.5,
"np.float64": 64.5,
"np.int16": -16,
"np.int32": -32,
"np.int64": -64,
"np.int8": -8,
"np.uint16": 16,
"np.uint32": 32,
"np.uint64": 64,
"np.uint8": 8,
}
] == list(fresh_db["types"].rows)