From 7ebfad1d40c03a426cf2b582e72247351a62eb68 Mon Sep 17 00:00:00 2001
From: Nuno Saavedra <32967989+Nfsaavedra@users.noreply.github.com>
Date: Tue, 23 Jan 2024 13:57:55 +0000
Subject: [PATCH] Set size of library cache (#38)

---
 coqpyt/coq/proof_file.py   | 16 ++++++++++++++++
 coqpyt/tests/test_cache.py | 10 ++++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 coqpyt/tests/test_cache.py

diff --git a/coqpyt/coq/proof_file.py b/coqpyt/coq/proof_file.py
index 03cc50a..ad46942 100644
--- a/coqpyt/coq/proof_file.py
+++ b/coqpyt/coq/proof_file.py
@@ -174,6 +174,12 @@ def __load_library(
         coq_file.close()
         return context
 
+    @staticmethod
+    def set_cache_size(size: Optional[int]):
+        _AuxFile._AuxFile__load_library = lru_cache(maxsize=size)(
+            _AuxFile.__load_library.__wrapped__,
+        )
+
     @staticmethod
     def get_library(
         library_name: str,
@@ -732,6 +738,16 @@ def __is_proven(self, proof: ProofTerm) -> bool:
             and "Admitted" not in proof.steps[-1].step.short_text
         )
 
+    @staticmethod
+    def set_library_cache_size(size: Optional[int]):
+        """Sets the size of the cache used to store the libraries of the Coq files.
+
+        Args:
+            size (Optional[int]): The size of the cache. If None, the cache
+                will have no limit.
+        """
+        _AuxFile.set_cache_size(size)
+
     @property
     def proofs(self) -> List[ProofTerm]:
         """Gets all the closed proofs in the file and their corresponding steps.
diff --git a/coqpyt/tests/test_cache.py b/coqpyt/tests/test_cache.py
new file mode 100644
index 0000000..890f968
--- /dev/null
+++ b/coqpyt/tests/test_cache.py
@@ -0,0 +1,10 @@
+from coqpyt.coq.proof_file import _AuxFile, ProofFile
+
+
+def test_set_cache_size():
+    _AuxFile.set_cache_size(256)
+    _AuxFile._AuxFile__load_library.cache_info().maxsize == 256
+    _AuxFile.set_cache_size(512)
+    _AuxFile._AuxFile__load_library.cache_info().maxsize == 512
+    ProofFile.set_library_cache_size(256)
+    _AuxFile._AuxFile__load_library.cache_info().maxsize == 256