Merge pull request #52 from FabioRosado/headers

Allows the first argument of headers to be a str or list

tests/test_decorators.py

@@ -110,6 +110,15 @@ def test_headers(request_builder):
     headers.modify_request(request_builder)
     assert request_builder.info["headers"] == {"key_1": "value_1"}
 
+    headers_str = decorators.headers("key_1: value_1")
+    headers_str.modify_request(request_builder)
+    assert request_builder.info["headers"] == {"key_1": "value_1"}
+
+    headers_lst = decorators.headers(["key_1: value_1", "key_2: value_2"])
+    headers_lst.modify_request(request_builder)
+    header = {"key_1": "value_1", "key_2": "value_2"}
+    assert request_builder.info["headers"] == header
+
 
 def test_form_url_encoded(request_builder):
     form_url_encoded = decorators.form_url_encoded()

uplink/decorators.py

@@ -122,10 +122,16 @@ class GitHub(Consumer):
     """
 
     def __init__(self, arg, **kwargs):
-        # TODO: allow the first argument to be a list or str, in
-        # which case you would split the strings by a colon delimiter
-        # to identify the headers.
-        self._headers = dict(arg, **kwargs)
+        if isinstance(arg, str):
+            key, value = self._get_header(arg)
+            self._headers = {key: value}
+        elif isinstance(arg, list):
+            self._headers = dict(self._get_header(a) for a in arg)
+        else:
+            self._headers = dict(arg, **kwargs)
+
+    def _get_header(self, arg):
+        return map(str.strip, arg.split(":"))
 
     def modify_request(self, request_builder):
         """Updates header contents."""