add tests + refactoring

Author: Dacry

dumpconverter.py

@@ -13,5 +13,5 @@
     #parser.add_argument("-p", "--processes", help="Number of processors to use (default 4)", type=int, default=4)
     args = parser.parse_args()
     start = time.time()
-    CsvWriter.write_csv(newJsonReader.process_json(args.input), args.output)
+    CsvWriter.write_csv(JsonReader.read_json(args.input), args.output)
     print "total time: %.2fs" % (time.time() - start)

propertysuggester/parser/JsonReader.py

@@ -1,9 +1,9 @@
 """
-process_json returns a generator that yields Entities)
+read_json returns a generator that yields Entities)
 
 usage:
 with open("file.csv", "r") as f:
-    for entity in process_json(f):
+    for entity in read_json(f):
         do_things()
 
 """
@@ -15,7 +15,7 @@
     print "ujson not found"
     import json as json
 
-def process_json(input_file):
+def read_json(input_file):
     count = 0
     for jsonline in input_file:
         count += 1 
@@ -27,33 +27,34 @@ def process_json(input_file):
         except ValueError:
             continue
         if data["type"] == "item":
-            title = data["id"]
-            if not "claims" in data:
-                yield Entity(title, [])
-                continue
-            claims = []
-            for prop, statements in data["claims"].iteritems():
-                for statement in statements:
-                    references = []
-                    if "references" in statement:
-                        for prop, snaks in statement["references"][0]["snaks"].iteritems():
-                            for snak in snaks:
-                                ref = _parse_json_snak(snak)
-                                if ref:
-                                    references.append(ref)
-                    qualifiers = []
-                    if "qualifiers" in statement:                            
-                        for prop, snaks in statement["qualifiers"].iteritems():
-                            for snak in snaks: 
-                                qualifier = _parse_json_snak(snak)
-                            if qualifier:
-                                qualifiers.append(qualifier)
-                    claim = _parse_json_snak(statement["mainsnak"])
-                    if claim:
-                        claims.append(Claim(claim, qualifiers, references))
+            yield _process_json(data)
 
-            yield Entity(title, claims)
+def _process_json(data):
+    title = data["id"]
+    if not "claims" in data:
+        return Entity(title, [])
+    claims = []
+    for prop, statements in data["claims"].iteritems():
+        for statement in statements:
+            references = []
+            if "references" in statement:
+                for prop, snaks in statement["references"][0]["snaks"].iteritems():
+                    for snak in snaks:
+                        ref = _parse_json_snak(snak)
+                        if ref:
+                            references.append(ref)
+            qualifiers = []
+            if "qualifiers" in statement:                            
+                for prop, snaks in statement["qualifiers"].iteritems():
+                    for snak in snaks: 
+                        qualifier = _parse_json_snak(snak)
+                    if qualifier:
+                        qualifiers.append(qualifier)
+            claim = _parse_json_snak(statement["mainsnak"])
+            if claim:
+                claims.append(Claim(claim, qualifiers, references))
 
+    return Entity(title, claims)
 
 def _parse_json_snak(claim_json):
     if claim_json["snaktype"] == "value":
@@ -81,5 +82,5 @@ def _parse_json_snak(claim_json):
     else:  # novalue, somevalue, ...
         datatype = "unknown"
         value = claim_json["snaktype"]
-    property_id = claim_json["property"][1:]
+    property_id = int(claim_json["property"][1:])
     return Snak(property_id, datatype, value)

propertysuggester/parser/XmlReader.py

@@ -114,6 +114,7 @@ def _parse_json_snak(claim_json):
         if datatype == "string":
             value = claim_json[3]
         elif datatype == "wikibase-entityid":
+            datatype = "wikibase-item"
             if claim_json[3]["entity-type"] == "item":
                 value = "Q" + str(claim_json[3]["numeric-id"])
             else:

propertysuggester/test/parser/Wikidata-Q1.json.gz

@@ -11,8 +11,8 @@ def assert_universe(self, result):
 
         self.assertThat(q1.title, Equals("Q1"))
         self.assertThat(q1.claims, Contains(Claim(Snak(373, "string", "Universe"), [],
-                                                  [Snak(143, "wikibase-entityid", "Q328")])))
-        self.assertThat(q1.claims, Contains(Claim(Snak(31, "wikibase-entityid", "Q223557"))))
-        self.assertThat(q1.claims, Contains(Claim(Snak(31, "wikibase-entityid", "Q1088088"))))
-        self.assertThat(q1.claims, Contains(Claim(Snak(361, "wikibase-entityid", "Q3327819"),
-                                                  [Snak(31, "wikibase-entityid", "Q41719")], [])))
+                                                  [Snak(143, "wikibase-item", "Q328")])))
+        self.assertThat(q1.claims, Contains(Claim(Snak(31, "wikibase-item", "Q223557"))))
+        self.assertThat(q1.claims, Contains(Claim(Snak(31, "wikibase-item", "Q1088088"))))
+        self.assertThat(q1.claims, Contains(Claim(Snak(361, "wikibase-item", "Q3327819"),
+                                                  [Snak(31, "wikibase-item", "Q41719")], [])))

propertysuggester/test/parser/Wikidata-Q15511.json.gz

@@ -16,11 +16,11 @@ def setUp(self):
     def test_universe(self):
         out = StringIO()
         out.writelines(["Q1,claim,373,string,Universe\n",
-                        "Q1,reference,143,wikibase-entityid,Q328\n"
-                        "Q1,claim,31,wikibase-entityid,Q223557\n",
-                        "Q1,claim,31,wikibase-entityid,Q1088088\n",
-                        "Q1,claim,361,wikibase-entityid,Q3327819\n",
-                        "Q1,qualifier,31,wikibase-entityid,Q41719\n"])
+                        "Q1,reference,143,wikibase-item,Q328\n"
+                        "Q1,claim,31,wikibase-item,Q223557\n",
+                        "Q1,claim,31,wikibase-item,Q1088088\n",
+                        "Q1,claim,361,wikibase-item,Q3327819\n",
+                        "Q1,qualifier,31,wikibase-item,Q41719\n"])
         out.seek(0)
         result = list(CsvReader.read_csv(out))
         self.assert_universe(result)

propertysuggester/test/parser/test_abstract_reader.py

@@ -0,0 +1,30 @@
+import unittest
+import gzip
+
+from pkg_resources import resource_filename
+from testtools import TestCase
+from testtools.matchers import *
+
+from propertysuggester.test.parser.test_abstract_reader import AbstractUniverseTest
+from propertysuggester.parser import JsonReader
+from propertysuggester.utils.datamodel import Claim, Snak, Entity
+
+
+class JsonReaderTest(AbstractUniverseTest):
+
+    def test_updated_dump(self):
+        with gzip.open(resource_filename(__name__, "Wikidata-Q15511.json.gz"), "r") as f:
+            result = list(JsonReader.read_json(f))
+
+        self.assertThat(len(result), Equals(1))
+        q15511 = result[0]
+        self.assertThat(q15511.title, Equals("Q15511"))
+        self.assertThat(q15511.claims, Contains(Claim(Snak(1082, "quantity", "+25"), [Snak(585, "time", "+00000002001-01-01T00:00:00Z"), Snak(459, "wikibase-item", "Q745221")], [Snak(248, "wikibase-item", "Q17597573")])))
+
+    def test_special_cases(self):
+        data = dict([("id", "Q1"), ("type", "item")])
+        self.assertThat(JsonReader._process_json(data), Equals(Entity("Q1", [])))
+
+if __name__ == '__main__':
+    unittest.main()
+    

propertysuggester/test/parser/test_csv_reader.py

@@ -23,7 +23,7 @@ def test_updated_dump(self):
         self.assertThat(len(result), Equals(1))
         q9351 = result[0]
         self.assertThat(q9351.title, Equals("Q9351"))
-        self.assertThat(q9351.claims, Contains(Claim(Snak(156, "wikibase-entityid", "Q1647331"))))
+        self.assertThat(q9351.claims, Contains(Claim(Snak(156, "wikibase-item", "Q1647331"))))
         self.assertThat(q9351.claims, Contains(Claim(Snak(1112, "quantity", "+25"))))
 
     def test_special_cases(self):