Added typing to improve the use of the library (#1)

benchmarks/conftest.py

@@ -1,46 +1,51 @@
-import os
 import math
+import os
+from typing import Final
+
 import pytest
 
-from pymerkle import SqliteTree as MerkleTree, constants
+from pymerkle import SqliteTree as MerkleTree
+from pymerkle import constants
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
+current_dir: str = os.path.dirname(os.path.abspath(__file__))
 
-DEFAULT_DB = os.path.join(current_dir, 'merkle.db')
-DEFAULT_SIZE = 10 ** 6
-DEFAULT_INDEX = math.ceil(DEFAULT_SIZE / 2)
-DEFAULT_ROUNDS = 100
-DEFAULT_THRESHOLD = 128
-DEFAULT_CAPACITY = 1024 ** 3
+DEFAULT_DB: Final[str] = os.path.join(current_dir, 'merkle.db')
+DEFAULT_SIZE: Final[int] = 10 ** 6
+DEFAULT_INDEX: Final[int] = math.ceil(DEFAULT_SIZE / 2)
+DEFAULT_ROUNDS: Final[int] = 100
+DEFAULT_THRESHOLD: Final[int] = 128
+DEFAULT_CAPACITY: Final[int] = 1024 ** 3
 
 
 def pytest_addoption(parser):
     parser.addoption('--dbfile', type=str, default=DEFAULT_DB,
-        help='Database filepath')
+                     help='Database filepath')
     parser.addoption('--size', type=int, default=DEFAULT_SIZE,
-        help='Nr entries to consider')
+                     help='Nr entries to consider')
     parser.addoption('--index', type=int, default=DEFAULT_INDEX,
-        help='Base index for proof operations')
+                     help='Base index for proof operations')
     parser.addoption('--rounds', type=int, default=DEFAULT_ROUNDS,
-        help='Nr rounds per benchmark')
+                     help='Nr rounds per benchmark')
     parser.addoption('--algorithm', default='sha256',
-        choices=constants.ALGORITHMS,
-        help='Hash algorithm used by the tree')
+                     choices=constants.ALGORITHMS,
+                     help='Hash algorithm used by the tree')
     parser.addoption('--randomize', action='store_true', default=False,
-        help='Randomize function input per round')
+                     help='Randomize function input per round')
     parser.addoption('--disable-optimizations', action='store_true', default=False,
-        help='Use unoptimized versions of core operations')
+                     help='Use unoptimized versions of core operations')
     parser.addoption('--disable-cache', action='store_true', default=False,
-        help='Disable subroot caching')
+                     help='Disable subroot caching')
     parser.addoption('--threshold', type=int, metavar='WIDTH',
-        default=DEFAULT_THRESHOLD,
-        help='Subroot cache threshold')
+                     default=DEFAULT_THRESHOLD,
+                     help='Subroot cache threshold')
     parser.addoption('--capacity', type=int, metavar='BYTES',
-        default=DEFAULT_CAPACITY,
-        help='Subroot cache capacity in bytes')
+                     default=DEFAULT_CAPACITY,
+                     help='Subroot cache capacity in bytes')
+
 
 option = None
 
-def pytest_configure(config):
+
+def pytest_configure(config) -> None:
     global option
     option = config.option

benchmarks/init_db.py

@@ -3,77 +3,80 @@
 the database file will be overwritten if it already exists.
 """
 
+import argparse
 import os
 import sys
-import argparse
 import time
+from typing import Any, Final
 
 from pymerkle import SqliteTree, constants
 
-current_dir = os.path.dirname(os.path.abspath(__file__))
+current_dir: str = os.path.dirname(os.path.abspath(path=__file__))
 
-DEFAULT_DB = os.path.join(current_dir, 'merkle.db')
-DEFAULT_ALGORITHM = 'sha256'
-DEFAULT_SIZE = 10 ** 8
-DEFAULT_BATCHSIZE = 10 ** 7
+DEFAULT_DB: Final[str] = os.path.join(current_dir, 'merkle.db')
+DEFAULT_ALGORITHM: Final[str] = 'sha256'
+DEFAULT_SIZE: Final[int] = 10 ** 8
+DEFAULT_BATCHSIZE: Final[int] = 10 ** 7
 
 
-def parse_cli_args():
-    config = {'prog': sys.argv[0], 'usage': 'python %s' % sys.argv[0],
-              'description': __doc__, 'epilog': '\n',
-              'formatter_class': argparse.ArgumentDefaultsHelpFormatter}
+def parse_cli_args() -> argparse.Namespace:
+    config: dict[str, Any] = {'prog': sys.argv[0], 'usage': 'python %s' % sys.argv[0],
+                              'description': __doc__, 'epilog': '\n',
+                              'formatter_class': argparse.ArgumentDefaultsHelpFormatter}
     parser = argparse.ArgumentParser(**config)
 
     parser.add_argument('--dbfile', type=str, default=DEFAULT_DB,
-        help='Database filepath')
+                        help='Database filepath')
     parser.add_argument('--algorithm', choices=constants.ALGORITHMS,
-        default=DEFAULT_ALGORITHM, help='Hashing algorithm')
+                        default=DEFAULT_ALGORITHM, help='Hashing algorithm')
     parser.add_argument('--disable-security', action='store_true', default=False,
-        help='Disable resistance against 2nd-preimage attack')
+                        help='Disable resistance against 2nd-preimage attack')
     parser.add_argument('--size', type=int, default=DEFAULT_SIZE,
-        help='Nr entries to append in total')
+                        help='Nr entries to append in total')
     parser.add_argument('--batchsize', type=int, default=DEFAULT_BATCHSIZE,
-        help='Nr entries to append per bulk insertion')
+                        help='Nr entries to append per bulk insertion')
     parser.add_argument('--preserve-database', action='store_true', default=False,
-        help='Append without overwriting if already existent')
+                        help='Append without overwriting if already existent')
 
     return parser.parse_args()
 
 
 if __name__ == '__main__':
-    args = parse_cli_args()
+    args: argparse.Namespace = parse_cli_args()
 
-    batchsize = args.batchsize
+    batchsize = int(args.batchsize)
     size = args.size
     if batchsize > size:
         sys.stdout.write("[-] Batchsize exceeds size\n")
         sys.exit(1)
 
     if not args.preserve_database:
         try:
-            os.remove(args.dbfile)
+            os.remove(path=args.dbfile)
         except OSError:
             pass
 
-    opts = {'algorithm': args.algorithm,
-            'disable_security': args.disable_security}
+    opts: dict[str, Any] = {'algorithm': args.algorithm,
+                            'disable_security': args.disable_security}
 
-    with SqliteTree(args.dbfile, **opts) as tree:
-        offset = 0
-        count = 1
+    with SqliteTree(dbfile=args.dbfile, **opts) as tree:
+        offset: int = 0
+        count: int = 1
         append_entries = tree.append_entries
-        chunksize = min(100_000, batchsize)
-        start_time = time.time()
+        chunksize: int = min(100_000, batchsize)
+        start_time: float = time.time()
+        currsize: int = 0
         while offset < size:
-            limit = offset + batchsize + 1
+            limit: int = offset + batchsize + 1
             if limit > size + 1:
                 limit = size + 1
 
             print(f"\nCreating {batchsize} entries...")
-            entries = [f'entry-{i}'.encode('utf-8') for i in range(offset + 1,
-                limit)]
+            entries: list[bytes] = [f'entry-{i}'.encode(encoding='utf-8') for i in range(offset + 1,
+                                                                                         limit)]
 
-            index = append_entries(entries, chunksize)
+            index: int = append_entries(
+                entries=entries, chunksize=chunksize)  # type: ignore
             assert index == limit - 1
 
             currsize = tree.get_size()
@@ -83,8 +86,8 @@ def parse_cli_args():
             count += 1
             offset += batchsize
 
-        end_time = time.time()
-        elapsed_time = end_time - start_time
+        end_time: float = time.time()
+        elapsed_time: float = end_time - start_time
 
         assert currsize == args.size
 

benchmarks/test_perf.py

@@ -1,21 +1,24 @@
 from random import randint
+from typing import Any
+
 import pytest
 
 from pymerkle import SqliteTree as MerkleTree
+
 from .conftest import option
 
-defaults = {'warmup_rounds': 0, 'rounds': option.rounds}
+defaults: dict[str, Any] = {'warmup_rounds': 0, 'rounds': option.rounds}
 
 
-opts = {'disable_optimizations': option.disable_optimizations,
-        'disable_cache': option.disable_cache,
-        'threshold': option.threshold,
-        'capacity': option.capacity}
+opts: dict[str, Any] = {'disable_optimizations': option.disable_optimizations,
+                        'disable_cache': option.disable_cache,
+                        'threshold': option.threshold,
+                        'capacity': option.capacity}
 
-tree = MerkleTree(option.dbfile, algorithm=option.algorithm, **opts)
+tree = MerkleTree(dbfile=option.dbfile, algorithm=option.algorithm, **opts)
 
 
-def test_root(benchmark):
+def test_root(benchmark) -> None:
 
     def setup():
         start = randint(0, option.size - 2) if option.randomize else 0
@@ -27,7 +30,7 @@ def setup():
     benchmark.pedantic(tree._get_root, setup=setup, **defaults)
 
 
-def test_state(benchmark):
+def test_state(benchmark) -> None:
 
     def setup():
         size = randint(1, option.size) if option.randomize \
@@ -38,7 +41,7 @@ def setup():
     benchmark.pedantic(tree.get_state, setup=setup, **defaults)
 
 
-def test_inclusion(benchmark):
+def test_inclusion(benchmark) -> None:
 
     def setup():
         size = option.size
@@ -49,7 +52,7 @@ def setup():
     benchmark.pedantic(tree.prove_inclusion, setup=setup, **defaults)
 
 
-def test_consistency(benchmark):
+def test_consistency(benchmark) -> None:
 
     def setup():
         size2 = option.size

demo.py

@@ -29,19 +29,19 @@ def parse_cli_args():
     parser = argparse.ArgumentParser(**config)
 
     parser.add_argument('--backend', choices=['inmemory', 'sqlite'],
-            default='inmemory', help='Storage backend')
+                        default='inmemory', help='Storage backend')
     parser.add_argument('--algorithm', choices=constants.ALGORITHMS,
-            default='sha256', help='Hashing algorithm')
+                        default='sha256', help='Hashing algorithm')
     parser.add_argument('--threshold', type=int, metavar='WIDTH',
-            default=128, help='Subroot cache threshold')
+                        default=128, help='Subroot cache threshold')
     parser.add_argument('--capacity', type=int, metavar='MAXSIZE',
-            default=1024 ** 3, help='Subroot cache capacity in bytes')
+                        default=1024 ** 3, help='Subroot cache capacity in bytes')
     parser.add_argument('--disable-security', action='store_true',
-            default=False, help='Disable resistance against second-preimage attack')
+                        default=False, help='Disable resistance against second-preimage attack')
     parser.add_argument('--disable-optimizations', action='store_true',
-            default=False, help='Use unopmitized versions of core operations')
+                        default=False, help='Use unopmitized versions of core operations')
     parser.add_argument('--disable-cache', action='store_true',
-            default=False, help='Disable subroot caching')
+                        default=False, help='Disable subroot caching')
 
     return parser.parse_args()
 
@@ -50,7 +50,7 @@ def order_of_magnitude(num):
     return int(log10(num)) if not num == 0 else 0
 
 
-def strpath(rule, path):
+def strpath(rule, path) -> str:
     s2 = 3 * ' '
     s3 = 3 * ' '
     template = '\n{s1}[{index}]{s2}{bit}{s3}{value}'
@@ -65,16 +65,16 @@ def strpath(rule, path):
     return ''.join(pairs)
 
 
-def strtree(tree):
+def strtree(tree) -> str:
     if isinstance(tree, SqliteTree):
         entries = [tree.get_entry(index) for index in range(1, tree.get_size()
-            + 1)]
+                                                            + 1)]
         tree = InmemoryTree.init_from_entries(entries)
 
     return str(tree)
 
 
-def strproof(proof):
+def strproof(proof) -> str:
     template = """
     algorithm   : {algorithm}
     security    : {security}
@@ -98,7 +98,7 @@ def strproof(proof):
 if __name__ == '__main__':
     args = parse_cli_args()
 
-    MerkleTree = { 'inmemory': InmemoryTree, 'sqlite': SqliteTree }[
+    MerkleTree = {'inmemory': InmemoryTree, 'sqlite': SqliteTree}[
         args.backend]
 
     config = {'algorithm': args.algorithm,