Add function for senses and forms

Update senses.py

Author: LeMyst

wikibaseintegrator/entities/lexeme.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
 import re
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from wikibaseintegrator.entities.baseentity import BaseEntity
-from wikibaseintegrator.models.forms import Forms
+from wikibaseintegrator.models.forms import Form, Forms
 from wikibaseintegrator.models.lemmas import Lemmas
-from wikibaseintegrator.models.senses import Senses
+from wikibaseintegrator.models.senses import Sense, Senses
 from wikibaseintegrator.wbi_config import config
+from wikibaseintegrator.wbi_helpers import lexeme_add_form, lexeme_add_sense
 
 
 class LexemeEntity(BaseEntity):
@@ -163,3 +164,27 @@ def write(self, **kwargs: Any) -> LexemeEntity:
         """
         json_data = super()._write(data=self.get_json(), **kwargs)
         return self.from_json(json_data=json_data)
+
+    def write_form(self, form: Form) -> str:
+        if not self.id:
+            raise Exception('You must set a Lexeme id before writing a Form.')
+        return lexeme_add_form(lexeme_id=self.id, data=form.get_json())['form']['id']
+
+    def write_forms(self) -> List[str]:
+        ids: List = []
+        for form in self.forms:
+            ids.append(self.write_form(form))
+
+        return ids
+
+    def write_sense(self, sense: Sense) -> str:
+        if not self.id:
+            raise Exception('You must set a Lexeme id before writing a Sense.')
+        return lexeme_add_sense(lexeme_id=self.id, data=sense.get_json())['sense']['id']
+
+    def write_senses(self) -> List[str]:
+        ids: List = []
+        for sense in self.senses:
+            ids.append(self.write_sense(sense))
+
+        return ids

wikibaseintegrator/models/forms.py

@@ -40,6 +40,9 @@ def get_json(self) -> List[Dict]:
 
         return json_data
 
+    def __iter__(self):
+        return self.forms
+
     def __len__(self):
         return len(self.forms)
 

wikibaseintegrator/models/senses.py

@@ -37,6 +37,9 @@ def get_json(self) -> List[Dict]:
 
         return json_data
 
+    def __iter__(self):
+        return self.senses
+
     def __len__(self):
         return len(self.senses)
 
@@ -75,6 +78,9 @@ def remove(self) -> Sense:
         self.removed = True
         return self
 
+    def __eq__(self, other):
+        return self.glosses == other.glosses and self.claims == other.glosses
+
 
 class Glosses(LanguageValues):
     pass