From fae944292f64463a2ce6a2a660e225e6211979d0 Mon Sep 17 00:00:00 2001
From: Myst <1592048+LeMyst@users.noreply.github.com>
Date: Tue, 27 Dec 2022 21:50:56 +0100
Subject: [PATCH] Add function for senses and forms

Update senses.py
---
 wikibaseintegrator/entities/lexeme.py | 29 +++++++++++++++++++++++++--
 wikibaseintegrator/models/forms.py    |  3 +++
 wikibaseintegrator/models/senses.py   |  6 ++++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/wikibaseintegrator/entities/lexeme.py b/wikibaseintegrator/entities/lexeme.py
index 1c18e1e4..cadd878d 100644
--- a/wikibaseintegrator/entities/lexeme.py
+++ b/wikibaseintegrator/entities/lexeme.py
@@ -4,10 +4,11 @@
 from typing import Any
 
 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
diff --git a/wikibaseintegrator/models/forms.py b/wikibaseintegrator/models/forms.py
index bd8b5d99..3021054b 100644
--- a/wikibaseintegrator/models/forms.py
+++ b/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)
 
diff --git a/wikibaseintegrator/models/senses.py b/wikibaseintegrator/models/senses.py
index 4778fa78..fd49816e 100644
--- a/wikibaseintegrator/models/senses.py
+++ b/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