rebased onto latest in master branch of python-gitlab/python-gitlab

d0c-s4vage · d0c-s4vage · commit 8141f7757237 · 2017-08-22T12:29:09.000-05:00
diff --git a/gitlab/__init__.py b/gitlab/__init__.py
@@ -396,11 +396,11 @@ def _raw_list(self, path_, cls, **kwargs):
 
         return results
 
-    def _raw_post(self, path_, data=None, content_type=None, **kwargs):
+    def _raw_post(self, path_, data=None, content_type=None, files=None, **kwargs):
         url = '%s%s' % (self._url, path_)
         opts = self._get_session_opts(content_type)
         try:
-            return self.session.post(url, params=kwargs, data=data, **opts)
+            return self.session.post(url, params=kwargs, data=data, files=files, **opts)
         except Exception as e:
             raise GitlabConnectionError(
                 "Can't connect to GitLab server (%s)" % e)
diff --git a/gitlab/exceptions.py b/gitlab/exceptions.py
@@ -176,6 +176,13 @@ class GitlabTimeTrackingError(GitlabOperationError):
 class GitlabCherryPickError(GitlabOperationError):
     pass
 
+class GitlabUploadError(GitlabOperationError):
+    pass
+
+
+class GitlabAttachFileError(GitlabOperationError):
+    pass
+
 
 def raise_error_from_response(response, error, expected_code=200):
     """Tries to parse gitlab error message from response and raises error.
diff --git a/gitlab/v3/objects.py b/gitlab/v3/objects.py
@@ -66,6 +66,55 @@ def compound_metrics(self, **kwargs):
         return self._simple_get('/sidekiq/compound_metrics', **kwargs)
 
 
+class ProjectUploadable(object):
+    """A mixin for objects that allow files to be uploaded/attached.
+    """
+
+    description_attr = "description"
+    """The attribute of the object that newly uploaded files'
+    markdown should be appended to.
+    """
+
+    project_id_attr = "project_id"
+    """The attribute that specifies the attribute that contains
+    the project id into which files may be uploaded.
+    """
+
+    # see #56 - add file uploading/attachment features
+    def attach_file(self, filename, filedata, **kwargs):
+        """Attach a file to the issue
+        
+        Args:
+            filename (str): The filename of the file being uploaded.
+            filedata (str): The raw data of the file being uploaded.
+        Raises:
+            GitlabConnectionError: If the server cannot be reached.
+        """
+        project_id = getattr(self, self.project_id_attr, None)
+        if project_id is None:
+            raise GitlabAttachFileError("{}'s project id could not be determined (tried using {!r})".format(
+                self,
+                self.project_id_attr,
+            ))
+
+        project = self.gitlab.projects.get(project_id)
+        res = project.upload(
+            filename = filename,
+            filedata = filedata,
+        )
+
+        orig_desc = getattr(self, self.description_attr, "")
+        # files are "attached" to issues and comments by uploading the
+        # the file into the project, and including the returned
+        # markdown in the description
+        setattr(self, self.description_attr, orig_desc + "\n\n" + res["markdown"])
+
+        # XXX:TODO: Is this correct? any changes to the current object instance 
+        # that have not yet been saved will be saved along with the
+        # file upload, which *may not* be what the user intended.
+        self.save()
+
+
 class UserEmail(GitlabObject):
     _url = '/users/%(user_id)s/emails'
     canUpdate = False
@@ -901,20 +950,24 @@ class ProjectHookManager(BaseManager):
     obj_cls = ProjectHook
 
 
-class ProjectIssueNote(GitlabObject):
+class ProjectIssueNote(GitlabObject, ProjectUploadable):
     _url = '/projects/%(project_id)s/issues/%(issue_id)s/notes'
     _constructorTypes = {'author': 'User'}
     canDelete = False
     requiredUrlAttrs = ['project_id', 'issue_id']
     requiredCreateAttrs = ['body']
     optionalCreateAttrs = ['created_at']
 
+    # file attachment settings (see #56)
+    description_attr = "body"
+    project_id_attr = "project_id"
+
 
 class ProjectIssueNoteManager(BaseManager):
     obj_cls = ProjectIssueNote
 
 
-class ProjectIssue(GitlabObject):
+class ProjectIssue(GitlabObject, ProjectUploadable):
     _url = '/projects/%(project_id)s/issues/'
     _constructorTypes = {'author': 'User', 'assignee': 'User',
                          'milestone': 'ProjectMilestone'}
@@ -933,6 +986,10 @@ class ProjectIssue(GitlabObject):
             [('project_id', 'project_id'), ('issue_id', 'id')]),
     )
 
+    # file attachment settings (see #56)
+    description_attr = "description"
+    project_id_attr = "project_id"
+
     def subscribe(self, **kwargs):
         """Subscribe to an issue.
 
@@ -1057,6 +1114,7 @@ class ProjectIssueManager(BaseManager):
 
 class ProjectMember(GitlabObject):
     _url = '/projects/%(project_id)s/members'
+
     requiredUrlAttrs = ['project_id']
     requiredCreateAttrs = ['access_level', 'user_id']
     optionalCreateAttrs = ['expires_at']
@@ -2096,6 +2154,39 @@ def trigger_build(self, ref, token, variables={}, **kwargs):
         r = self.gitlab._raw_post(url, data=data, **kwargs)
         raise_error_from_response(r, GitlabCreateError, 201)
 
+    # see #56 - add file attachment features
+    def upload(self, filename, filedata, **kwargs):
+        """Upload a file into the project. This will return the raw response
+        from the gitlab API in the form:
+        
+        {
+          "alt": "dk",
+          "url": "/uploads/66dbcd21ec5d24ed6ea225176098d52b/dk.png",
+          "markdown": "![dk](/uploads/66dbcd21ec5d24ed6ea225176098d52b/dk.png)"
+        }
+
+        See https://github.com/gitlabhq/gitlabhq/blob/master/doc/api/projects.md#upload-a-file
+        for more information.
+
+        Args:
+            filename (str): The name of the file being uploaded
+            filedata (bytes): The raw data of the file being uploaded
+
+        Raises:
+            GitlabConnectionError: If the server cannot be reached
+            GitlabUploadError: If the file upload fails
+        """
+        url = ("/projects/%(id)s/uploads" % {
+            "id": self.id,
+        })
+        r = self.gitlab._raw_post(
+            url,
+            files = {"file": (filename, filedata)},
+        )
+        # returns 201 status code (created)
+        raise_error_from_response(r, GitlabUploadError, expected_code=201)
+        return r.json()
+
 
 class Runner(GitlabObject):
     _url = '/runners'
