diff --git a/docs/python-api.rst b/docs/python-api.rst index ebfc203..35023d5 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -1136,6 +1136,27 @@ This produces the following schema: [LatinName] TEXT ) +You can use the ``rename=`` argument to rename columns in the lookup table. To create a ``Species`` table with columns called ``name`` and ``latin`` you can do this: + +.. code-block:: python + + db["Trees"].extract( + ["CommonName", "LatinName"], + table="Species", + fk_column="species_id", + rename={"CommonName": "name", "LatinName": "latin"} + ) + +This produces a lookup table like so: + +.. code-block:: sql + + CREATE TABLE [Species] ( + [id] INTEGER PRIMARY KEY, + [name] TEXT, + [latin] TEXT + ) + .. _python_api_hash: Setting an ID based on the hash of the row contents diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 846f97d..33af372 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -883,7 +883,8 @@ class Table(Queryable): return sqls - def extract(self, columns, table=None, fk_column=None): + def extract(self, columns, table=None, fk_column=None, rename=None): + rename = rename or {} if isinstance(columns, str): columns = [columns] if not set(columns).issubset(self.columns_dict.keys()): @@ -898,7 +899,7 @@ class Table(Queryable): lookup_table = self.db[table] for row in self.rows: row_pks = tuple(row[pk] for pk in pks) - lookups = {column: row[column] for column in columns} + lookups = {rename.get(column) or column: row[column] for column in columns} self.update(row_pks, {first_column: lookup_table.lookup(lookups)}) fk_column = fk_column or "{}_id".format(table) # Now rename first_column and change its type to integer, and drop diff --git a/tests/test_extract.py b/tests/test_extract.py index 31b6d3e..4855441 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -45,7 +45,7 @@ def test_extract_single_column(fresh_db, table, fk_column): ] -def test_extract_multiple_columns(fresh_db): +def test_extract_multiple_columns_with_rename(fresh_db): iter_common = itertools.cycle(["Palm", "Spruce", "Mangrove", "Oak"]) iter_latin = itertools.cycle(["Arecaceae", "Picea", "Rhizophora", "Quercus"]) fresh_db["tree"].insert_all( @@ -61,7 +61,9 @@ def test_extract_multiple_columns(fresh_db): pk="id", ) - fresh_db["tree"].extract(["common_name", "latin_name"]) + fresh_db["tree"].extract(["common_name", "latin_name"], rename={ + "common_name": "name" + }) assert fresh_db["tree"].schema == ( 'CREATE TABLE "tree" (\n' " [id] INTEGER PRIMARY KEY,\n" @@ -73,15 +75,15 @@ def test_extract_multiple_columns(fresh_db): assert fresh_db["common_name_latin_name"].schema == ( "CREATE TABLE [common_name_latin_name] (\n" " [id] INTEGER PRIMARY KEY,\n" - " [common_name] TEXT,\n" + " [name] TEXT,\n" " [latin_name] TEXT\n" ")" ) assert list(fresh_db["common_name_latin_name"].rows) == [ - {"common_name": "Palm", "id": 1, "latin_name": "Arecaceae"}, - {"common_name": "Spruce", "id": 2, "latin_name": "Picea"}, - {"common_name": "Mangrove", "id": 3, "latin_name": "Rhizophora"}, - {"common_name": "Oak", "id": 4, "latin_name": "Quercus"}, + {"name": "Palm", "id": 1, "latin_name": "Arecaceae"}, + {"name": "Spruce", "id": 2, "latin_name": "Picea"}, + {"name": "Mangrove", "id": 3, "latin_name": "Rhizophora"}, + {"name": "Oak", "id": 4, "latin_name": "Quercus"}, ] assert list(itertools.islice(fresh_db["tree"].rows, 0, 4)) == [ {"id": 1, "name": "Tree 1", "common_name_latin_name_id": 1},