can use operators on update queries

Author: dantownsend

docs/src/piccolo/query_types/update.rst

@@ -14,6 +14,36 @@ This is used to update any rows in the table which match the criteria.
     >>> ).run_sync()
     []
 
+You can also add / subtract / multiply / divide values:
+
+.. code-block:: python
+
+    # Add 100 to the popularity of each band:
+    Band.update({
+        Band.popularity: Band.popularity + 100
+    }).run_sync()
+
+    # Decrease the popularity of each band by 100.
+    Band.update({
+        Band.popularity: Band.popularity - 100
+    }).run_sync()
+
+    # Multiply the popularity of each band by 10.
+    Band.update({
+        Band.popularity: Band.popularity * 10
+    }).run_sync()
+
+    # Divide the popularity of each band by 10.
+    Band.update({
+        Band.popularity: Band.popularity / 10
+    }).run_sync()
+
+    # You can also use the operators in reverse:
+    Band.update({
+        Band.popularity: 2000 - Band.popularity
+    }).run_sync()
+
+
 Query clauses
 -------------
 

piccolo/columns/column_types.py

@@ -5,7 +5,7 @@
 import uuid
 
 from piccolo.columns.base import Column, OnDelete, OnUpdate, ForeignKeyMeta
-from piccolo.querystring import Unquoted
+from piccolo.querystring import Unquoted, QueryString
 
 if t.TYPE_CHECKING:
     from piccolo.table import Table  # noqa
@@ -82,6 +82,49 @@ def __init__(self, default: int = None, **kwargs) -> None:
         kwargs.update({"default": default})
         super().__init__(**kwargs)
 
+    def _get_querystring(self, operator: str, value: int, reverse=False):
+        """
+        Used in update queries - for example:
+
+        await Band.update({Band.popularity: Band.popularity + 100}).run()
+        """
+        if not isinstance(value, (int, float)):
+            raise ValueError("Only integers and floats can be added.")
+        if reverse:
+            return QueryString(f"{{}} {operator} {self._meta.name} ", value)
+        else:
+            return QueryString(f"{self._meta.name} {operator} {{}}", value)
+
+    def __add__(self, value: int):
+        return self._get_querystring("+", value)
+
+    def __radd__(self, value: int):
+        return self._get_querystring("+", value, reverse=True)
+
+    def __sub__(self, value: int):
+        return self._get_querystring("-", value)
+
+    def __rsub__(self, value: int):
+        return self._get_querystring("-", value, reverse=True)
+
+    def __mul__(self, value: int):
+        return self._get_querystring("*", value)
+
+    def __rmul__(self, value: int):
+        return self._get_querystring("*", value, reverse=True)
+
+    def __truediv__(self, value: int):
+        return self._get_querystring("/", value)
+
+    def __rtruediv__(self, value: int):
+        return self._get_querystring("/", value, reverse=True)
+
+    def __floordiv__(self, value: int):
+        return self._get_querystring("/", value)
+
+    def __rfloordiv__(self, value: int):
+        return self._get_querystring("/", value, reverse=True)
+
 
 ###############################################################################
 # BigInt and SmallInt only exist on Postgres. SQLite treats them the same as

tests/table/test_update.py

@@ -6,7 +6,7 @@ class TestUpdate(DBTestCase):
     def test_update(self):
         self.insert_rows()
 
-        Band.update().values({Band.name: "Pythonistas3"}).where(
+        Band.update({Band.name: "Pythonistas3"}).where(
             Band.name == "Pythonistas"
         ).run_sync()
 
@@ -18,3 +18,93 @@ def test_update(self):
         print(f"response = {response}")
 
         self.assertEqual(response, [{"name": "Pythonistas3"}])
+
+    def test_update_values(self):
+        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.assertEqual(response, [{"name": "Pythonistas3"}])
+
+
+class TestUpdateOperators(DBTestCase):
+    def test_add(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: Band.popularity + 10}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 1010)
+
+    def test_radd(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: 10 + Band.popularity}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 1010)
+
+    def test_sub(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: Band.popularity - 10}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 990)
+
+    def test_rsub(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: 1100 - Band.popularity}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 100)
+
+    def test_mul(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: Band.popularity * 2}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 2000)
+
+    def test_rmul(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: 2 * Band.popularity}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 2000)
+
+    def test_div(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: Band.popularity / 10}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 100)
+
+    def test_rdiv(self):
+        self.insert_row()
+
+        Band.update({Band.popularity: 1000 / Band.popularity}).run_sync()
+
+        response = Band.select(Band.popularity).first().run_sync()
+
+        self.assertEqual(response["popularity"], 1)