From 0b2a1ff6421657b5bdff234962e1f467e8cd4bbc Mon Sep 17 00:00:00 2001
From: zmezei <mezeienator@gmail.com>
Date: Thu, 6 Jul 2023 23:16:42 +0100
Subject: [PATCH] refactor, avoid multiple json load and dump

---
 src/uagents/models.py     | 25 ++++++++++++++-----------
 src/uagents/protocol.py   |  2 +-
 tests/test_field_descr.py |  6 +++---
 3 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/src/uagents/models.py b/src/uagents/models.py
index 12d6dd3d..703e8d9d 100644
--- a/src/uagents/models.py
+++ b/src/uagents/models.py
@@ -1,21 +1,22 @@
 import hashlib
 from typing import Type, Union, Dict, Any
 import json
+import copy
 from pydantic import BaseModel
 
 
 class Model(BaseModel):
-    _schema_no_descr = None
+    _schema_no_descriptions = None
 
     @staticmethod
     def remove_descriptions(schema: Dict[str, Dict[str, str]]):
-        fields_with_descr = []
         if not "properties" in schema:
             return
+
+        fields_with_descr = []
         for field_name, field_props in schema["properties"].items():
             if "description" in field_props:
                 fields_with_descr.append(field_name)
-
         for field_name in fields_with_descr:
             del schema["properties"][field_name]["description"]
 
@@ -24,21 +25,23 @@ def remove_descriptions(schema: Dict[str, Dict[str, str]]):
                 Model.remove_descriptions(definition)
 
     @classmethod
-    def schema_no_descr(cls) -> Dict[str, Any]:
-        if cls._schema_no_descr is None:
-            schema = json.loads(cls.schema_json(indent=None, sort_keys=True))
+    def schema_no_descriptions(cls) -> Dict[str, Any]:
+        if cls._schema_no_descriptions is None:
+            schema = copy.deepcopy(cls.schema())
             Model.remove_descriptions(schema)
-            cls._schema_no_descr = schema
-        return cls._schema_no_descr
+            cls._schema_no_descriptions = schema
+        return cls._schema_no_descriptions
 
     @classmethod
-    def schema_json_no_descr(cls) -> str:
-        return json.dumps(cls.schema_no_descr())
+    def schema_json_no_descriptions(cls) -> str:
+        return json.dumps(cls.schema_no_descriptions(), indent=None, sort_keys=True)
 
     @staticmethod
     def build_schema_digest(model: Union["Model", Type["Model"]]) -> str:
         digest = (
-            hashlib.sha256(model.schema_json_no_descr().encode("utf-8")).digest().hex()
+            hashlib.sha256(model.schema_json_no_descriptions().encode("utf-8"))
+            .digest()
+            .hex()
         )
         return f"model:{digest}"
 
diff --git a/src/uagents/protocol.py b/src/uagents/protocol.py
index 8fa93069..6ffb26e8 100644
--- a/src/uagents/protocol.py
+++ b/src/uagents/protocol.py
@@ -176,7 +176,7 @@ def manifest(self) -> Dict[str, Any]:
 
         for schema_digest, model in all_models.items():
             manifest["models"].append(
-                {"digest": schema_digest, "schema": model.schema_no_descr()}
+                {"digest": schema_digest, "schema": model.schema_no_descriptions()}
             )
 
         for request, responses in self._replies.items():
diff --git a/tests/test_field_descr.py b/tests/test_field_descr.py
index 671b680f..38d4270e 100644
--- a/tests/test_field_descr.py
+++ b/tests/test_field_descr.py
@@ -41,7 +41,7 @@ class Message(Model):
 
         self.assertEqual(
             Message.schema_json(indent=None, sort_keys=True),
-            Message.schema_json_no_descr(),
+            Message.schema_json_no_descriptions(),
         )
 
         class Message(Model):
@@ -50,7 +50,7 @@ class Message(Model):
 
         self.assertNotEqual(
             Message.schema_json(indent=None, sort_keys=True),
-            Message.schema_json_no_descr(),
+            Message.schema_json_no_descriptions(),
         )
 
         class MessageArgs(Model):
@@ -63,7 +63,7 @@ class Message(Model):
 
         self.assertNotEqual(
             Message.schema_json(indent=None, sort_keys=True),
-            Message.schema_json_no_descr(),
+            Message.schema_json_no_descriptions(),
         )
 
     def test_model_digest(self):