implement document and field set

2 files changed, 106 insertions, 20 deletions

diff --git a/python/vespa/vespa/package.py b/python/vespa/vespa/package.py
index 988eb724906..e8f9dd49d83 100644
--- a/python/vespa/vespa/package.py
+++ b/python/vespa/vespa/package.py
@@ -5,24 +5,6 @@ import docker
 from vespa.json_serialization import ToJson, FromJson
 
 
-class Document(object):
-    def __init__(self) -> None:
-        """
-        Object representing a Vespa document.
-
-        """
-        self.fields = []
-
-    def add_fields(self, *fields):
-        """
-        Add Fields to the document.
-
-        :param fields: fields to be added
-        :return:
-        """
-        self.fields.extend(list(fields))
-
-
 class Field(ToJson, FromJson["Field"]):
     def __init__(
         self,
@@ -76,6 +58,73 @@ class Field(ToJson, FromJson["Field"]):
         return "{0}\n{1}".format(self.__class__.__name__, str(self.to_dict))
 
 
+class Document(ToJson, FromJson["Document"]):
+    def __init__(self, fields: Optional[List[Field]] = None) -> None:
+        """
+        Object representing a Vespa document.
+
+        """
+        if not fields:
+            fields = []
+
+        self.fields = fields
+
+    def add_fields(self, *fields: Field):
+        """
+        Add Fields to the document.
+
+        :param fields: fields to be added
+        :return:
+        """
+        self.fields.extend(list(fields))
+
+    @staticmethod
+    def from_dict(mapping: Mapping) -> "Document":
+        return Document(fields=[FromJson.map(field) for field in mapping.get("fields")])
+
+    @property
+    def to_dict(self) -> Mapping:
+        map = {"fields": [field.to_envelope for field in self.fields]}
+        return map
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return self.fields == other.fields
+
+    def __repr__(self):
+        return "{0}\n{1}".format(self.__class__.__name__, str(self.to_dict))
+
+
+class FieldSet(ToJson, FromJson["FieldSet"]):
+    def __init__(self, name: str, fields: List[str]) -> None:
+        """
+        A fieldset groups fields together for searching.
+
+        :param name: Name of the fieldset
+        :param fields: Field names to be included in the fieldset.
+        """
+        self.name = name
+        self.fields = fields
+
+    @staticmethod
+    def from_dict(mapping: Mapping) -> "FieldSet":
+        return FieldSet(name=mapping["name"], fields=mapping["fields"])
+
+    @property
+    def to_dict(self) -> Mapping:
+        map = {"name": self.name, "fields": self.fields}
+        return map
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return self.name == other.name and self.fields == other.fields
+
+    def __repr__(self):
+        return "{0}\n{1}".format(self.__class__.__name__, str(self.to_dict))
+
+
 class ApplicationPackage(object):
     def __init__(self, name: str, disk_folder: str) -> None:
         """
diff --git a/python/vespa/vespa/test_package.py b/python/vespa/vespa/test_package.py
index 281a30bebf0..6da0db1eafd 100644
--- a/python/vespa/vespa/test_package.py
+++ b/python/vespa/vespa/test_package.py
@@ -1,6 +1,6 @@
 import unittest
 
-from vespa.package import Field
+from vespa.package import Field, Document, FieldSet
 
 
 class TestField(unittest.TestCase):
@@ -42,4 +42,41 @@ class TestField(unittest.TestCase):
             ),
         )
         self.assertEqual(field, Field.from_dict(field.to_dict))
-        print(str(field))
+
+
+class TestDocument(unittest.TestCase):
+    def test_empty_document(self):
+        document = Document()
+        self.assertEqual(document.fields, [])
+        self.assertEqual(document.to_dict, {"fields": []})
+        self.assertEqual(document, Document.from_dict(document.to_dict))
+
+    def test_document_one_field(self):
+        document = Document()
+        field = Field(name="test_name", type="string")
+        document.add_fields(field)
+        self.assertEqual(document.fields, [field])
+        self.assertEqual(document, Document.from_dict(document.to_dict))
+        self.assertEqual(document, Document([field]))
+
+    def test_document_two_fields(self):
+        document = Document()
+        field_1 = Field(name="test_name", type="string")
+        field_2 = Field(
+            name="body",
+            type="string",
+            indexing=["index", "summary"],
+            index=["enable-bm25"],
+        )
+        document.add_fields(field_1, field_2)
+        self.assertEqual(document.fields, [field_1, field_2])
+        self.assertEqual(document, Document.from_dict(document.to_dict))
+        self.assertEqual(document, Document([field_1, field_2]))
+
+
+class TestFieldSet(unittest.TestCase):
+    def test_fieldset(self):
+        field_set = FieldSet(name="default", fields=["title", "body"])
+        self.assertEqual(field_set.name, "default")
+        self.assertEqual(field_set.fields, ["title", "body"])
+        self.assertEqual(field_set, FieldSet.from_dict(field_set.to_dict))