diff --git a/src/knowledge/data/chunk.py b/src/knowledge/data/chunk.py index 5267dc3..76f2058 100644 --- a/src/knowledge/data/chunk.py +++ b/src/knowledge/data/chunk.py @@ -1,9 +1,9 @@ -from enum import Enum +from enum import IntEnum from ..object import ObjectID ChunkID = ObjectID -class PositionType(Enum): +class PositionType(IntEnum): Unknown = 1 Device = 2 File = 3 diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index b3841ed..490555b 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -1,4 +1,3 @@ - # define a object type enum from abc import ABC, abstractmethod from enum import Enum @@ -6,6 +5,16 @@ from .object_id import ObjectID, ObjectType import hashlib import json import pickle +from typing import Any + + +class ObjectEnhancedJSONEncoder(json.JSONEncoder): + def default(self, o: Any) -> Any: + if isinstance(o, ObjectID): + return o.to_base58() + + return super().default(o) + class KnowledgeObject(ABC): def __init__(self, object_type: ObjectType, desc: dict = {}, body: dict = {}): @@ -18,7 +27,7 @@ class KnowledgeObject(ABC): def object_id(self) -> ObjectID: return self.calculate_id() - + def set_desc_with_key_value(self, key, value): self.desc[key] = value @@ -27,7 +36,7 @@ class KnowledgeObject(ABC): def get_desc(self) -> dict: return self.desc - + def set_body_with_key_value(self, key, value): self.body[key] = value @@ -36,10 +45,13 @@ class KnowledgeObject(ABC): def get_body(self) -> dict: return self.body - + def calculate_id(self): # Convert the object_type and desc to string and compute the SHA256 hash - data = json.dumps({"object_type": self.object_type, "desc": self.desc}) + data = json.dumps( + {"object_type": self.object_type, "desc": self.desc}, + cls=ObjectEnhancedJSONEncoder, + ) sha256 = hashlib.sha256() sha256.update(data.encode()) return ObjectID(sha256.digest()) @@ -50,5 +62,3 @@ class KnowledgeObject(ABC): @staticmethod def decode(data: bytes): return pickle.loads(data) - - diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 8278183..42f313a 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -1,12 +1,12 @@ # define a object type enum from abc import ABC, abstractmethod -from enum import Enum +from enum import IntEnum from .hash import HashValue import base58 import base36 -class ObjectType(Enum): +class ObjectType(IntEnum): Chunk = 7 Image = 101 Video = 102 diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py index 34c0911..f273279 100644 --- a/src/knowledge/object/relation.py +++ b/src/knowledge/object/relation.py @@ -4,10 +4,10 @@ import sqlite3 from typing import List, Tuple, Optional import logging import os -from enum import Enum +from enum import IntEnum -class ObjectRelationType(Enum): +class ObjectRelationType(IntEnum): Parent = 1 diff --git a/test/test_object.py b/test/test_object.py index 0f5620e..655e2a2 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -23,6 +23,7 @@ from knowledge import ( EmailObjectBuilder, ObjectRelationStore, KnowledgeStore, + EmailObject, ) import asyncio import unittest @@ -42,6 +43,20 @@ class TestVectorSTorage(unittest.TestCase): email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" email_object = EmailObjectBuilder({}, email_folder).build() + + id = email_object.calculate_id() + print(f"got email object: {id.to_base58()}") + + # test encode & decode + ret = email_object.encode() + obj = EmailObject.decode(ret) + id2 = obj.calculate_id() + print(f"got email object: {id2.to_base58()}") + self.assertEqual(id.to_base58(), id2.to_base58()) + + ret2 = obj.encode() + self.assertEqual(ret, ret2) + def test_relation(self): obj1 = ObjectID.hash_data("12345".encode("utf-8"))