Add function for senses and forms

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):
@@ -147,3 +148,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,10 +40,12 @@ def get_json(self) -> List[Dict]:
 
         return json_data
 
+    def __iter__(self):
+        return self.forms
+
     def __len__(self):
         return len(self.forms)
 
-
 class Form(BaseModel):
     def __init__(self, form_id: Optional[str] = None, representations: Optional[Representations] = None, grammatical_features: Optional[Union[str, int, List[str]]] = None, claims: Optional[Claims] = None):
         self.id = form_id