From e615d22c5564ea1f32abb086088b2700110be10a Mon Sep 17 00:00:00 2001
From: Simon Willison <simonw@eventbrite.com>
Date: Sat, 23 Feb 2019 20:02:19 -0800
Subject: [PATCH] Support for numpy types, closes #11 (#12)

---
 .travis.yml          |  2 ++
 sqlite_utils/db.py   | 22 +++++++++++++
 tests/test_create.py | 78 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 102 insertions(+)

diff --git a/.travis.yml b/.travis.yml
index 6dd695f..3af29a6 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -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:
diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py
index 35f2fc2..83d54a4 100644
--- a/sqlite_utils/db.py
+++ b/sqlite_utils/db.py
@@ -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,
diff --git a/tests/test_create.py b/tests/test_create.py
index 4bd0af5..7bf9ad8 100644
--- a/tests/test_create.py
+++ b/tests/test_create.py
@@ -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)