From aa81f3b2303a60e27dd34477f66ba35097a61464 Mon Sep 17 00:00:00 2001
From: Myst <1592048+LeMyst@users.noreply.github.com>
Date: Fri, 23 Jun 2023 23:48:09 +0200
Subject: [PATCH] Implementation EntityShape

Fix #470
---
 requirements.txt                          |  1 +
 setup.cfg                                 |  1 +
 setup.py                                  |  1 +
 test/test_entity_item.py                  |  7 +++++++
 wikibaseintegrator/entities/baseentity.py | 21 +++++++++++++++++++++
 5 files changed, 31 insertions(+)

diff --git a/requirements.txt b/requirements.txt
index 5da49924..f8bc9814 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,3 +4,4 @@ oauthlib~=3.2.2
 requests~=2.31.0
 requests-oauthlib~=1.3.1
 ujson~=5.9.0
+entityshape~=0.1.0
diff --git a/setup.cfg b/setup.cfg
index cb904dd2..8a465ad6 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -40,6 +40,7 @@ install_requires =
     requests>=2.27.1,<2.29.0
     requests-oauthlib~=1.3.1
     ujson>=5.4,<5.6
+    entityshape~=0.1.0
 python_requires = >=3.8, <3.13
 
 [options.extras_require]
diff --git a/setup.py b/setup.py
index b7e9bd4f..7d7ced60 100644
--- a/setup.py
+++ b/setup.py
@@ -11,6 +11,7 @@
         "requests >= 2.27.1,< 2.32.0",
         "requests-oauthlib ~= 1.3.1",
         "ujson >= 5.4,< 5.10"
+        "entityshape ~= 0.1.0"
     ],
     extras_require={
         "dev": [
diff --git a/test/test_entity_item.py b/test/test_entity_item.py
index a83ed3aa..92a49cbf 100644
--- a/test/test_entity_item.py
+++ b/test/test_entity_item.py
@@ -107,3 +107,10 @@ def test_new_lines(self):
             item.claims.add(MonolingualText(prop_nr=123, text="Multi\r\nline"))
             item.claims.add(MonolingualText(prop_nr=123, text="Multi\rline"))
             item.claims.add(MonolingualText(prop_nr=123, text="Multi\nline"))
+
+    def test_entity_schema(self):
+        random_campsite = wbi.item.get('Q119156070')
+        assert random_campsite.schema_validator(entity_schema_id="E376").is_valid
+        assert random_campsite.schema_validator(entity_schema_id="376").is_valid
+        assert random_campsite.schema_validator(entity_schema_id=376).is_valid
+        assert not wbi.item.get('Q582').schema_validator(entity_schema_id="E376").is_valid
diff --git a/wikibaseintegrator/entities/baseentity.py b/wikibaseintegrator/entities/baseentity.py
index b7857b55..3d12bcab 100644
--- a/wikibaseintegrator/entities/baseentity.py
+++ b/wikibaseintegrator/entities/baseentity.py
@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 import logging
+import re
 from copy import copy
 from typing import TYPE_CHECKING, Any
 
+from entityshape import EntityShape, Result
+
 from wikibaseintegrator import wbi_fastrun
 from wikibaseintegrator.datatypes import BaseDataType
 from wikibaseintegrator.models.claims import Claim, Claims
+from wikibaseintegrator.wbi_config import config
 from wikibaseintegrator.wbi_enums import ActionIfExists
 from wikibaseintegrator.wbi_exceptions import MissingEntityException
 from wikibaseintegrator.wbi_helpers import delete_page, edit_entity, mediawiki_api_call_helper
@@ -306,6 +310,23 @@ def download_entity_ttl(self, **kwargs) -> str:
 
         raise ValueError('entity ID is null')
 
+    def schema_validator(self, entity_schema_id: str, language: str | None = None) -> Result:
+        if isinstance(entity_schema_id, str):
+            pattern = re.compile(r'^(?:[a-zA-Z]+:)?E?([0-9]+)$')
+            matches = pattern.match(entity_schema_id)
+
+            if not matches:
+                raise ValueError(f"Invalid EntitySchema ID ({entity_schema_id}), format must be 'E[0-9]+'")
+
+            entity_schema_id = f'E{matches.group(1)}'
+        elif isinstance(entity_schema_id, int):
+            entity_schema_id = f'E{entity_schema_id}'
+        else:
+            raise ValueError(f"Invalid EntitySchema ID ({entity_schema_id}), format must be 'E[0-9]+'")
+
+        language = str(language or config['DEFAULT_LANGUAGE'])
+        return EntityShape(qid=self.id, eid=entity_schema_id, lang=language).validate_and_get_result()
+
     def __repr__(self):
         """A mixin implementing a simple __repr__."""
         return "<{klass} @{id:x} {attrs}>".format(  # pylint: disable=consider-using-f-string