improved table updating

Dictionaries with string keys, and kwargs can now be used.

Author: dantownsend

piccolo/query/methods/update.py

@@ -19,10 +19,13 @@ class Update(Query):
 
     def __init__(self, table: t.Type[Table], **kwargs):
         super().__init__(table, **kwargs)
-        self.values_delegate = ValuesDelegate()
+        self.values_delegate = ValuesDelegate(table=table)
         self.where_delegate = WhereDelegate()
 
-    def values(self, values: t.Dict[Column, t.Any]) -> Update:
+    def values(
+        self, values: t.Dict[t.Union[Column, str], t.Any] = {}, **kwargs
+    ) -> Update:
+        values = dict(values, **kwargs)
         self.values_delegate.values(values)
         return self
 

piccolo/query/mixins.py

@@ -234,16 +234,40 @@ def remove_secret_columns(self):
 class ValuesDelegate:
     """
     Used to specify new column values - primarily used in update queries.
-
-    Example usage:
-
-    .values({MyTable.column_a: 1})
     """
 
+    table: t.Type[Table]
     _values: t.Dict[Column, t.Any] = field(default_factory=dict)
 
-    def values(self, values: t.Dict[Column, t.Any]):
-        self._values.update(values)
+    def values(self, values: t.Dict[t.Union[Column, str], t.Any]):
+        """
+        Example usage:
+
+        .. code-block:: python
+
+            .values({MyTable.column_a: 1})
+
+            # Or:
+            .values({'column_a': 1})
+
+            # Or:
+            .values(column_a=1})
+
+        """
+        cleaned_values: t.Dict[Column, t.Any] = {}
+        for key, value in values.items():
+            if isinstance(key, Column):
+                column = key
+            elif isinstance(key, str):
+                column = self.table._meta.get_column_by_name(key)
+            else:
+                raise ValueError(
+                    f"Unrecognised key - {key} is neither a Column or the "
+                    "name of a Column."
+                )
+            cleaned_values[column] = value
+
+        self._values.update(cleaned_values)
 
     def get_sql_values(self) -> t.List[t.Any]:
         """

piccolo/table.py

@@ -626,7 +626,9 @@ def table_exists(cls) -> TableExists:
         return TableExists(table=cls)
 
     @classmethod
-    def update(cls, values: t.Dict[Column, t.Any] = {}) -> Update:
+    def update(
+        cls, values: t.Dict[t.Union[Column, str], t.Any] = {}, **kwargs
+    ) -> Update:
         """
         Update rows.
 

tests/table/test_update.py

@@ -3,37 +3,90 @@
 
 
 class TestUpdate(DBTestCase):
+    def check_response(self):
+        response = (
+            Band.select(Band.name)
+            .where(Band.name == "Pythonistas3")
+            .run_sync()
+        )
+
+        self.assertEqual(response, [{"name": "Pythonistas3"}])
+
     def test_update(self):
+        """
+        Make sure updating work, when passing the new values directly to the
+        `update` method.
+        """
         self.insert_rows()
 
         Band.update({Band.name: "Pythonistas3"}).where(
             Band.name == "Pythonistas"
         ).run_sync()
 
-        response = (
-            Band.select(Band.name)
-            .where(Band.name == "Pythonistas3")
-            .run_sync()
-        )
-        print(f"response = {response}")
+        self.check_response()
 
-        self.assertEqual(response, [{"name": "Pythonistas3"}])
+    def test_update_with_string_keys(self):
+        """
+        Make sure updating work, when passing a dictionary of values, which
+        uses column names as keys, instead of Column instances.
+        """
+        self.insert_rows()
+
+        Band.update({"name": "Pythonistas3"}).where(
+            Band.name == "Pythonistas"
+        ).run_sync()
+
+        self.check_response()
+
+    def test_update_with_kwargs(self):
+        """
+        Make sure updating work, when passing the new value via kwargs.
+        """
+        self.insert_rows()
+
+        Band.update(name="Pythonistas3").where(
+            Band.name == "Pythonistas"
+        ).run_sync()
+
+        self.check_response()
 
     def test_update_values(self):
+        """
+        Make sure updating work, when passing the new values via the `values`
+        method.
+        """
         self.insert_rows()
 
         Band.update().values({Band.name: "Pythonistas3"}).where(
             Band.name == "Pythonistas"
         ).run_sync()
 
-        response = (
-            Band.select(Band.name)
-            .where(Band.name == "Pythonistas3")
-            .run_sync()
-        )
-        print(f"response = {response}")
+        self.check_response()
 
-        self.assertEqual(response, [{"name": "Pythonistas3"}])
+    def test_update_values_with_string_keys(self):
+        """
+        Make sure updating work, when passing the new values via the `values`
+        method, using a column name as a dictionary key.
+        """
+        self.insert_rows()
+
+        Band.update().values({"name": "Pythonistas3"}).where(
+            Band.name == "Pythonistas"
+        ).run_sync()
+
+        self.check_response()
+
+    def test_update_values_with_kwargs(self):
+        """
+        Make sure updating work, when passing the new values via kwargs.
+        """
+        self.insert_rows()
+
+        Band.update().values(name="Pythonistas3").where(
+            Band.name == "Pythonistas"
+        ).run_sync()
+
+        self.check_response()
 
 
 class TestIntUpdateOperators(DBTestCase):