Add type annotations to _add_labelled_object

khard/carddav_object.py

@@ -26,7 +26,8 @@
 from . import address_book  # pylint: disable=unused-import # for type checking
 from . import helpers
 from .helpers.typing import (Date, ObjectType, PostAddress, StrList,
-    convert_to_vcard, list_to_string, string_to_date, string_to_list)
+    StrListDict, convert_to_vcard, list_to_string, string_to_date,
+    string_to_list)
 from .query import AnyQuery, Query
 
 
@@ -387,14 +388,15 @@ def _get_new_group(self, group_type: str = "") -> str:
                 return group_name
 
     def _add_labelled_object(
-            self, obj_type: str, user_input, name_groups: bool = False,
+            self, obj_type: str, user_input: Union[StrList, StrListDict],
+            name_groups: bool = False,
             allowed_object_type: ObjectType = ObjectType.str) -> None:
         """Add an object to the VCARD. If user_input is a dict, the object will
          be added to a group with an ABLABEL created from the key of the dict.
 
         :param obj_type: type of object to add to the VCARD.
-        :param user_input: Contents of the object to add. If a dict
-        :type user_input: str or list(str) or dict(str) or dict(list(str))
+        :param user_input: Contents of the object to add. If a dict, the key
+            will be added as label
         :param name_groups: (Optional) If True, use the obj_type in the
             group name for labelled objects.
         :param allowed_object_type: (Optional) set the accepted return type

khard/helpers/typing.py

@@ -15,6 +15,7 @@ class ObjectType(Enum):
 Date = Union[str, datetime]
 StrList = Union[str, List[str]]
 PostAddress = Dict[str, str]
+StrListDict = Union[Dict[str, str], Dict[str, List[str]]]
 
 
 def convert_to_vcard(name: str, value: StrList, constraint: ObjectType

test/test_vcard_wrapper.py

@@ -1,13 +1,15 @@
 """Tests for the VCardWrapper class from the carddav module."""
 # pylint: disable=missing-docstring
 
+import contextlib
 import datetime
 import unittest
 from typing import Union, List, Dict
 
 import vobject
 
 from khard.carddav_object import VCardWrapper
+from khard.helpers.typing import ObjectType
 
 from .helpers import vCard, TestVCardWrapper
 
@@ -525,3 +527,50 @@ def test_setting_fn_from_labelled_org(self):
         wrapper._delete_vcard_object("FN")
         wrapper._add_organisation({'Work': ['Test Inc']})
         self.assertEqual(wrapper.formatted_name, 'Test Inc')
+
+
+class AddLabelledObject(unittest.TestCase):
+
+    @contextlib.contextmanager
+    def assertTitle(self, expected):
+        wrapper = TestVCardWrapper()
+        yield wrapper
+        self.assertEqual(wrapper._get_multi_property("TITLE"), expected)
+
+    def test_add_a_string(self):
+        with self.assertTitle(["foo"]) as wrapper:
+            wrapper._add_labelled_object("title", "foo")
+
+    def test_add_several_strings(self):
+        with self.assertTitle(["bar", "foo"]) as wrapper:
+            wrapper._add_labelled_object("title", "foo")
+            wrapper._add_labelled_object("title", "bar")
+
+    def test_add_a_list_of_strings(self):
+        with self.assertTitle([["foo","bar"]]) as wrapper:
+            wrapper._add_labelled_object("title", ["foo", "bar"],
+                                         allowed_object_type=ObjectType.list)
+
+    def test_add_string_with_label(self):
+        with self.assertTitle([{"foo": "bar"}]) as wrapper:
+            wrapper._add_labelled_object("title", {"foo": "bar"})
+
+    def test_add_strings_with_same_label(self):
+        with self.assertTitle([{"foo": "bar"}, {"foo": "baz"}]) as wrapper:
+            wrapper._add_labelled_object("title", {"foo": "bar"})
+            wrapper._add_labelled_object("title", {"foo": "baz"})
+
+    def test_add_strings_with_different_label(self):
+        with self.assertTitle([{"baz": "qux"}, {"foo": "bar"}]) as wrapper:
+            wrapper._add_labelled_object("title", {"foo": "bar"})
+            wrapper._add_labelled_object("title", {"baz": "qux"})
+
+    def test_add_a_list_with_label(self):
+        with self.assertTitle([{"foo": ["bar", "baz"]}]) as wrapper:
+            wrapper._add_labelled_object("title", {"foo": ["bar", "baz"]},
+                                         allowed_object_type=ObjectType.list)
+
+    def test_add_two_labelled_objects_at_once_fails(self):
+        wrapper = TestVCardWrapper()
+        with self.assertRaises(ValueError):
+            wrapper._add_labelled_object("title", {"a": "b", "c": "d"})