Use WikibaseDatatype enum

LeMyst · Jun 15, 2022 · c1e08ca · dpriskorn · Jun 15, 2022 · LeMyst
1 parent 20a5f85
commit c1e08ca
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/test/test_all.py b/test/test_all.py
@@ -5,7 +5,7 @@
 from wikibaseintegrator.datatypes import BaseDataType, Item
 from wikibaseintegrator.entities import ItemEntity
 from wikibaseintegrator.wbi_config import config as wbi_config
-from wikibaseintegrator.wbi_enums import ActionIfExists
+from wikibaseintegrator.wbi_enums import ActionIfExists, WikibaseDatatype
 from wikibaseintegrator.wbi_fastrun import get_fastrun_container
 
 wbi_config['USER_AGENT'] = 'WikibaseIntegrator-pytest/1.0 (test_all.py)'
@@ -19,7 +19,7 @@ def test_quantity(self):
 
         dt_json = dt.get_json()
 
-        assert dt_json['mainsnak']['datatype'] == 'quantity'
+        assert dt_json['mainsnak']['datatype'] == WikibaseDatatype.QUANTITY.value
 
         value = dt_json['mainsnak']['datavalue']
 
@@ -40,7 +40,7 @@ def test_geoshape(self):
 
         dt_json = dt.get_json()
 
-        assert dt_json['mainsnak']['datatype'] == 'geo-shape'
+        assert dt_json['mainsnak']['datatype'] == WikibaseDatatype.GEOSHAPE.value
 
         value = dt_json['mainsnak']['datavalue']
 

diff --git a/test/test_entity_property.py b/test/test_entity_property.py
@@ -32,3 +32,6 @@ def test_get(self):
 
     def test_get_json(self):
         assert wbi.property.get('P50', mediawiki_api_url='https://commons.wikimedia.org/w/api.php').get_json()['labels']['fr']['value'] == 'auteur ou autrice'
+
+    def test_create_property(self):
+        wbi.property.new(datatype='wikibase-item')
diff --git a/wikibaseintegrator/entities/property.py b/wikibaseintegrator/entities/property.py
@@ -7,6 +7,7 @@
 from wikibaseintegrator.models.aliases import Aliases
 from wikibaseintegrator.models.descriptions import Descriptions
 from wikibaseintegrator.models.labels import Labels
+from wikibaseintegrator.wbi_enums import WikibaseDatatype
 
 
 class PropertyEntity(BaseEntity):
@@ -16,7 +17,10 @@ def __init__(self, datatype: str = None, labels: Labels = None, descriptions: De
         super().__init__(**kwargs)
 
         # Property specific
-        self.datatype = datatype
+        if datatype is not None:
+            self.datatype = WikibaseDatatype(datatype)
+        else:
+            self.datatype = datatype
 
         # Item, Property and MediaInfo specific
         self.labels: Labels = labels or Labels()