fix(base): allow persisting attributes when updating object

nejch · nejch · commit 796c700cf854 · 2021-05-02T13:12:11.000+02:00
diff --git a/docs/api-usage.rst b/docs/api-usage.rst
@@ -190,6 +190,37 @@ a project (the previous example used 2 API calls):
    project = gl.projects.get(1, lazy=True)  # no API call
    project.star()  # API call
 
+.. _persist_attributes:
+
+Persisting local attributes
+===========================
+
+When methods manipulate an existing object, such as with ``refresh()`` and ``save()``,
+the object will only have attributes that were returned by the server. In some cases,
+such as when the initial request fetches attributes that are needed later for additional
+processing, this may not be desired:
+
+.. code-block:: python
+
+   project = gl.projects.get(1, statistics=True)
+   project.statistics
+
+   project.refresh()
+   project.statistics # AttributeError
+
+To avoid this, pass ``persist_attributes=True`` to ``refresh()``/``save()`` calls:
+
+.. code-block:: python
+
+   project = gl.projects.get(1, statistics=True)
+   project.statistics
+
+   project.refresh(persist_attributes=True)
+   project.statistics
+
+The ``persist_attributes`` setting is itself persisted in the object and can be reused
+for later ``refresh()`` and ``save()`` calls.
+
 Pagination
 ==========
 
diff --git a/docs/faq.rst b/docs/faq.rst
@@ -31,3 +31,8 @@ How can I clone the repository of a project?
         print(project.attributes)  # displays all the attributes
         git_url = project.ssh_url_to_repo
         subprocess.call(['git', 'clone', git_url])
+
+I get an ``AttributeError`` when accessing attributes after ``save()`` or ``refresh()``.
+    You are most likely trying to access an attribute that was not returned
+    by the server on the second request. Use the ``persist_attributes=True``
+    argument to override this - see :ref:`persist_attributes`.
diff --git a/gitlab/base.py b/gitlab/base.py
@@ -45,6 +45,7 @@ class RESTObject(object):
     _attrs: Dict[str, Any]
     _module: ModuleType
     _parent_attrs: Dict[str, Any]
+    _persist_attrs: bool
     _short_print_attr: Optional[str] = None
     _updated_attrs: Dict[str, Any]
     manager: "RESTManager"
@@ -59,6 +60,7 @@ def __init__(self, manager: "RESTManager", attrs: Dict[str, Any]) -> None:
             }
         )
         self.__dict__["_parent_attrs"] = self.manager.parent_attrs
+        self.__dict__["_persist_attrs"] = False
         self._create_managers()
 
     def __getstate__(self) -> Dict[str, Any]:
@@ -153,7 +155,11 @@ def _create_managers(self) -> None:
 
     def _update_attrs(self, new_attrs: Dict[str, Any]) -> None:
         self.__dict__["_updated_attrs"] = {}
-        self.__dict__["_attrs"] = new_attrs
+
+        if self.__dict__["_persist_attrs"] is True:
+            self.__dict__["_attrs"].update(new_attrs)
+        else:
+            self.__dict__["_attrs"] = new_attrs
 
     def get_id(self):
         """Returns the id of the resource."""
diff --git a/gitlab/mixins.py b/gitlab/mixins.py
@@ -162,10 +162,12 @@ class RefreshMixin(_RestObjectBase):
     manager: base.RESTManager
 
     @exc.on_http_error(exc.GitlabGetError)
-    def refresh(self, **kwargs: Any) -> None:
+    def refresh(self, persist_attributes: bool = None, **kwargs: Any) -> None:
         """Refresh a single object from server.
 
         Args:
+            persist_attributes: Whether to keep existing local attributes that
+            were not fetched from the server on refresh
             **kwargs: Extra options to send to the server (e.g. sudo)
 
         Returns None (updates the object)
@@ -174,6 +176,9 @@ def refresh(self, **kwargs: Any) -> None:
             GitlabAuthenticationError: If authentication is not correct
             GitlabGetError: If the server cannot perform the request
         """
+        if persist_attributes is not None:
+            self.__dict__["_persist_attrs"] = persist_attributes
+
         if self._id_attr:
             path = "%s/%s" % (self.manager.path, self.id)
         else:
@@ -529,18 +534,23 @@ def _get_updated_data(self) -> Dict[str, Any]:
 
         return updated_data
 
-    def save(self, **kwargs: Any) -> None:
+    def save(self, persist_attributes: bool = None, **kwargs: Any) -> None:
         """Save the changes made to the object to the server.
 
         The object is updated to match what the server returns.
 
         Args:
+            persist_attributes: Whether to keep existing local attributes that
+            were not fetched from the server on save
             **kwargs: Extra options to send to the server (e.g. sudo)
 
         Raise:
             GitlabAuthenticationError: If authentication is not correct
             GitlabUpdateError: If the server cannot perform the request
         """
+        if persist_attributes is not None:
+            self.__dict__["_persist_attrs"] = persist_attributes
+
         updated_data = self._get_updated_data()
         # Nothing to update. Server fails if sent an empty dict.
         if not updated_data:
diff --git a/gitlab/tests/test_base.py b/gitlab/tests/test_base.py
@@ -128,11 +128,25 @@ def test_update_attrs(self, fake_manager):
         assert {"foo": "foo", "bar": "bar"} == obj._attrs
         assert {} == obj._updated_attrs
 
-    def test_update_attrs_deleted(self, fake_manager):
-        obj = FakeObject(fake_manager, {"foo": "foo", "bar": "bar"})
-        obj.bar = "baz"
+    @pytest.mark.parametrize(
+        "initial_attrs,persist_attrs,assigned_attr,expected_attrs",
+        [
+            ({"foo": "foo", "bar": "bar"}, None, "baz", {"foo": "foo"}),
+            ({"foo": "foo", "bar": "bar"}, False, "baz", {"foo": "foo"}),
+            ({"foo": "foo", "bar": "bar"}, True, "baz", {"foo": "foo", "bar": "baz"}),
+        ],
+    )
+    def test_update_attrs_deleted(
+        self, fake_manager, initial_attrs, persist_attrs, assigned_attr, expected_attrs
+    ):
+        obj = FakeObject(fake_manager, initial_attrs)
+        obj._attrs["bar"] = assigned_attr
+
+        if persist_attrs is not None:
+            obj.__dict__["_persist_attrs"] = persist_attrs
+
         obj._update_attrs({"foo": "foo"})
-        assert {"foo": "foo"} == obj._attrs
+        assert expected_attrs == obj._attrs
         assert {} == obj._updated_attrs
 
     def test_dir_unique(self, fake_manager):
diff --git a/tools/functional/api/test_projects.py b/tools/functional/api/test_projects.py
@@ -234,6 +234,24 @@ def test_project_stars(project):
     assert project.star_count == 0
 
 
+@pytest.mark.parametrize(
+    "refresh_kwargs,hasattr_before,hasattr_after",
+    [
+        ({}, True, False),
+        ({"persist_attributes": True}, True, True),
+        ({"persist_attributes": False}, True, False),
+    ],
+)
+def test_project_statistics_after_refresh(
+    gl, project, refresh_kwargs, hasattr_before, hasattr_after
+):
+    project = gl.projects.get(project.id, statistics=True)
+    assert hasattr(project, "statistics") == hasattr_before
+
+    project.refresh(**refresh_kwargs)
+    assert hasattr(project, "statistics") == hasattr_after
+
+
 def test_project_tags(project, project_file):
     tag = project.tags.create({"tag_name": "v1.0", "ref": "master"})
     assert len(project.tags.list()) == 1
diff --git a/tools/functional/api/test_users.py b/tools/functional/api/test_users.py
@@ -142,6 +142,24 @@ def test_user_custom_attributes(gl, user):
     assert len(user.customattributes.list()) == 0
 
 
+@pytest.mark.parametrize(
+    "save_kwargs,hasattr_before,hasattr_after",
+    [
+        ({}, True, False),
+        ({"persist_attributes": True}, True, True),
+        ({"persist_attributes": False}, True, False),
+    ],
+)
+def test_user_custom_attributes_after_save(
+    gl, user, save_kwargs, hasattr_before, hasattr_after
+):
+    user = gl.users.get(user.id, with_custom_attributes=True)
+    assert hasattr(user, "custom_attributes") == hasattr_before
+
+    user.save(**save_kwargs)
+    assert hasattr(user, "custom_attributes") == hasattr_after
+
+
 def test_user_impersonation_tokens(gl, user):
     token = user.impersonationtokens.create(
         {"name": "token1", "scopes": ["api", "read_user"]}
