From b6c3255a9f1f35fbcd951c8faf9874bb64cea711 Mon Sep 17 00:00:00 2001
From: svlandeg <sofie.vanlandeghem@gmail.com>
Date: Thu, 21 Mar 2019 13:32:21 +0100
Subject: [PATCH] Entity class

---
 spacy/kb.pxd |  8 ++++++++
 spacy/kb.pyx | 22 ++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/spacy/kb.pxd b/spacy/kb.pxd
index 54ee49a3f..4ae34bfa7 100644
--- a/spacy/kb.pxd
+++ b/spacy/kb.pxd
@@ -43,6 +43,14 @@ cdef struct _AliasC:
     vector[float] probs
 
 
+# TODO: document
+cdef class Entity:
+
+    cdef readonly KnowledgeBase kb
+    cdef hash_t entity_hash
+    cdef float confidence
+
+
 # TODO: document
 cdef class Candidate:
 
diff --git a/spacy/kb.pyx b/spacy/kb.pyx
index 52c8ad8f0..4776e9d34 100644
--- a/spacy/kb.pyx
+++ b/spacy/kb.pyx
@@ -3,6 +3,28 @@
 from spacy.errors import user_warning
 
 
+cdef class Entity:
+
+    def __init__(self, KnowledgeBase kb, entity_hash, confidence):
+        self.kb = kb
+        self.entity_hash = entity_hash
+        self.confidence = confidence
+
+    property kb_id_:
+        """RETURNS (unicode): ID of this entity in the KB"""
+        def __get__(self):
+            return self.kb.strings[self.entity_hash]
+
+    property kb_id:
+        """RETURNS (uint64): hash of the entity's KB ID"""
+        def __get__(self):
+            return self.entity_hash
+
+    property confidence:
+        def __get__(self):
+            return self.confidence
+
+
 cdef class Candidate:
 
     def __init__(self, KnowledgeBase kb, entity_hash, alias_hash, prior_prob):