From 773f2b6b20622bb986984a1c3161d5b3aaa1046b Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 20 Jun 2022 12:46:49 -0700
Subject: [PATCH] Documented TypeTracker, closes #445

---
 docs/python-api.rst   | 59 +++++++++++++++++++++++++++++++++++++++++++
 docs/reference.rst    |  8 ++++++
 sqlite_utils/utils.py | 33 ++++++++++++++++++++++--
 3 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/docs/python-api.rst b/docs/python-api.rst
index 1a8c989..2a87610 100644
--- a/docs/python-api.rst
+++ b/docs/python-api.rst
@@ -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
diff --git a/docs/reference.rst b/docs/reference.rst
index 6603d99..abd918a 100644
--- a/docs/reference.rst
+++ b/docs/reference.rst
@@ -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
diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py
index d2ccc5f..98ceffd 100644
--- a/sqlite_utils/utils.py
+++ b/sqlite_utils/utils.py
@@ -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()}