Documented TypeTracker, closes #445

This commit is contained in:
Simon Willison 2022-06-20 12:46:49 -07:00
commit 773f2b6b20
3 changed files with 98 additions and 2 deletions

View file

@ -2588,6 +2588,65 @@ If you need to reset to the original value after calling this function you can d
csv.field_size_limit(ORIGINAL_CSV_FIELD_SIZE_LIMIT)
.. _python_api_typetracker:
Detecting column types using TypeTracker
========================================
Sometimes you may find yourself working with data that lacks type information - data from a CSV file for example.
The ``TypeTracker`` class can be used to try to automatically identify the most likely types for data that is initially represented as strings.
Consider this example:
.. code-block:: python
import csv, io
csv_file = io.StringIO("id,name\n1,Cleo\n2,Cardi")
rows = list(csv.DictReader(csv_file))
# rows is now this:
# [{'id': '1', 'name': 'Cleo'}, {'id': '2', 'name': 'Cardi'}]
If we insert this data directly into a table we will get a schema that is entirely ``TEXT`` columns:
.. code-block:: python
from sqlite_utils import Database
db = Database(memory=True)
db["creatures"].insert_all(rows)
print(db.schema)
# Outputs:
# CREATE TABLE [creatures] (
# [id] TEXT,
# [name] TEXT
# );
We can detect the best column types using a ``TypeTracker`` instance:
.. code-block:: python
from sqlite_utils.utils import TypeTracker
tracker = TypeTracker()
db["creatures2"].insert_all(tracker.wrap(rows))
print(tracker.types)
# Outputs {'id': 'integer', 'name': 'text'}
We can then apply those types to our new table using the :ref:`table.transform() <python_api_transform>` method:
.. code-block:: python
db["creatures2"].transform(types=tracker.types)
print(db["creatures2"].schema)
# Outputs:
# CREATE TABLE [creatures2] (
# [id] INTEGER,
# [name] TEXT
# );
.. _python_api_gis:
SpatiaLite helpers

View file

@ -86,3 +86,11 @@ sqlite_utils.utils.rows_from_file
---------------------------------
.. autofunction:: sqlite_utils.utils.rows_from_file
.. _reference_utils_typetracker:
sqlite_utils.utils.TypeTracker
------------------------------
.. autoclass:: sqlite_utils.utils.TypeTracker
:members: wrap, types

View file

@ -316,10 +316,35 @@ def rows_from_file(
class TypeTracker:
"""
Wrap an iterator of dictionaries and keep track of which SQLite column
types are the most likely fit for each of their keys.
Example usage:
.. code-block:: python
from sqlite_utils.utils import TypeTracker
import sqlite_utils
db = sqlite_utils.Database(memory=True)
tracker = TypeTracker()
rows = [{"id": "1", "name": "Cleo", "id": "2", "name": "Cardi"}]
db["creatures"].insert_all(tracker.wrap(rows))
print(tracker.types)
# Outputs {'id': 'integer', 'name': 'text'}
db["creatures"].transform(types=tracker.types)
"""
def __init__(self):
self.trackers = {}
def wrap(self, iterator):
def wrap(self, iterator: Iterable[dict]) -> Iterable[dict]:
"""
Use this to loop through an existing iterator, tracking the column types
as part of the iteration.
:param iterator: The iterator to wrap
"""
for row in iterator:
for key, value in row.items():
tracker = self.trackers.setdefault(key, ValueTracker())
@ -327,7 +352,11 @@ class TypeTracker:
yield row
@property
def types(self):
def types(self) -> Dict[str, str]:
"""
A dictionary mapping column names to their detected types. This can be passed
to the ``db[table_name].transform(types=tracker.types)`` method.
"""
return {key: tracker.guessed_type for key, tracker in self.trackers.items()}