refactor: simplify UpdateMany and UpdateOne __await__ method

BeanieODM · Aug 31, 2023 · bb773b7 · bb773b7
1 parent ff947d2
commit bb773b7
Showing 1 changed file with 29 additions and 29 deletions.
diff --git a/beanie/odm/queries/update.py b/beanie/odm/queries/update.py
@@ -195,17 +195,17 @@ def __await__(
         update_result = yield from self._update().__await__()
         if self.upsert_insert_doc is None:
             return update_result
-        else:
-            if update_result is not None and update_result.matched_count == 0:
-                return (
-                    yield from self.document_model.insert_one(
-                        document=self.upsert_insert_doc,
-                        session=self.session,
-                        bulk_writer=self.bulk_writer,
-                    ).__await__()
-                )
-            else:
-                return update_result
+
+        if update_result is not None and update_result.matched_count == 0:
+            return (
+                yield from self.document_model.insert_one(
+                    document=self.upsert_insert_doc,
+                    session=self.session,
+                    bulk_writer=self.bulk_writer,
+                ).__await__()
+            )
+
+        return update_result
 
 
 class UpdateOne(UpdateQuery):
@@ -347,21 +347,21 @@ def __await__(
         update_result = yield from self._update().__await__()
         if self.upsert_insert_doc is None:
             return update_result
-        else:
-            if (
-                self.response_type == UpdateResponse.UPDATE_RESULT
-                and update_result is not None
-                and update_result.matched_count == 0
-            ) or (
-                self.response_type != UpdateResponse.UPDATE_RESULT
-                and update_result is None
-            ):
-                return (
-                    yield from self.document_model.insert_one(
-                        document=self.upsert_insert_doc,
-                        session=self.session,
-                        bulk_writer=self.bulk_writer,
-                    ).__await__()
-                )
-            else:
-                return update_result
+
+        if (
+            self.response_type == UpdateResponse.UPDATE_RESULT
+            and update_result is not None
+            and update_result.matched_count == 0
+        ) or (
+            self.response_type != UpdateResponse.UPDATE_RESULT
+            and update_result is None
+        ):
+            return (
+                yield from self.document_model.insert_one(
+                    document=self.upsert_insert_doc,
+                    session=self.session,
+                    bulk_writer=self.bulk_writer,
+                ).__await__()
+            )
+
+        return update_result