Fix some type issues.

3 files changed, 24 insertions, 13 deletions

diff --git a/pyecsca/ec/configuration.py b/pyecsca/ec/configuration.py
index 0fd6976..24f3a7f 100644
--- a/pyecsca/ec/configuration.py
+++ b/pyecsca/ec/configuration.py
@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from itertools import product
-from typing import Set, get_type_hints, Union, get_origin, get_args, Generator
+from typing import get_type_hints, Union, get_origin, get_args, Generator, FrozenSet
 
 from public import public
 
@@ -78,7 +78,7 @@ class Configuration(object):
     """An ECC implementation configuration."""
     model: CurveModel
     coords: CoordinateModel
-    formulas: Set[Formula]
+    formulas: FrozenSet[Formula]
     scalarmult: ScalarMultiplier
     hash_type: HashType
     mod_rand: RandomMod
@@ -107,9 +107,10 @@ def all_configurations(**kwargs) -> Generator[Configuration, Configuration, None
     :param kwargs: The configuration parameters to match.
     :return: A generator of the configurations
     """
+
     def is_optional(arg_type):
         return get_origin(arg_type) == Union and len(get_args(arg_type)) == 2 and \
-               get_args(arg_type)[1] == type(None)
+               get_args(arg_type)[1] == type(None)  # noqa
 
     def leaf_subclasses(cls):
         subs = cls.__subclasses__()
@@ -151,12 +152,12 @@ def all_configurations(**kwargs) -> Generator[Configuration, Configuration, None
                                    isinstance(formula, opt_type)] + [None]
                     else:
                         options = [None]  # TODO: anything here?
-                elif get_origin(required_type) == None and issubclass(required_type, Formula):
+                elif get_origin(required_type) is None and issubclass(required_type, Formula):
                     options = [formula for formula in coords_formulas if
                                isinstance(formula, required_type)]
-                elif get_origin(required_type) == None and issubclass(required_type, bool):
+                elif get_origin(required_type) is None and issubclass(required_type, bool):
                     options = [True, False]
-                elif get_origin(required_type) == None and issubclass(required_type,
+                elif get_origin(required_type) is None and issubclass(required_type,
                                                                       int) and name == "width":
                     options = [3, 5]
                 else:
diff --git a/pyecsca/ec/curves.py b/pyecsca/ec/curves.py
index bc8329f..aa5e146 100644
--- a/pyecsca/ec/curves.py
+++ b/pyecsca/ec/curves.py
@@ -1,5 +1,6 @@
 import json
 from os.path import join
+from typing import Dict, Union
 
 from pkg_resources import resource_listdir, resource_isdir, resource_stream
 from public import public
@@ -41,6 +42,7 @@ def get_params(category: str, name: str, coords: str, infty: bool = True) -> Dom
     if curve["field"]["type"] == "Binary":
         raise ValueError("Binary field curves are currently not supported.")
 
+    # Get model and param names
     model: CurveModel
     field = int(curve["field"]["p"], 16)
     order = int(curve["order"], 16)
@@ -59,29 +61,36 @@ def get_params(category: str, name: str, coords: str, infty: bool = True) -> Dom
         param_names = ["a", "d"]
     else:
         raise ValueError("Unknown curve model.")
+
+    # Check coordinate model name and assumptions
     if coords not in model.coordinates:
         raise ValueError("Coordinate model not supported for curve.")
     coord_model = model.coordinates[coords]
     params = {name: Mod(int(curve["params"][name], 16), field) for name in param_names}
     for assumption in coord_model.assumptions:
-        locals = {}
+        alocals: Dict[str, Union[Mod, int]] = {}
         compiled = compile(assumption, "", mode="exec")
-        exec(compiled, None, locals)
-        for param, value in locals.items():
+        exec(compiled, None, alocals)
+        for param, value in alocals.items():
             if params[param] != value:
                 raise ValueError(f"Coordinate model {coord_model} has an unsatisifed assumption on the {param} parameter (= {value}).")
+    # Construct the point at infinity
+    infinity: Point
     if infty:
         infinity = InfinityPoint(coord_model)
     else:
-        locals = {**params}
+        ilocals: Dict[str, Union[Mod, int]] = {**params}
         for line in coord_model.neutral:
             compiled = compile(line, "", mode="exec")
-            exec(compiled, None, locals)
+            exec(compiled, None, ilocals)
         infinity_coords = {}
         for coordinate in coord_model.variables:
-            if coordinate not in locals:
+            if coordinate not in ilocals:
                 raise ValueError(f"Coordinate model {coord_model} requires infty option.")
-            infinity_coords[coordinate] = Mod(locals[coordinate], field)
+            value = ilocals[coordinate]
+            if isinstance(value, int):
+                value = Mod(value, field)
+            infinity_coords[coordinate] = value
         infinity = Point(coord_model, **infinity_coords)
     elliptic_curve = EllipticCurve(model, coord_model, field, infinity, params)
     affine = Point(AffineCoordinateModel(model), x=Mod(int(curve["generator"]["x"], 16), field),
diff --git a/pyecsca/ec/mod.py b/pyecsca/ec/mod.py
index 2b16a07..1fa3f9d 100644
--- a/pyecsca/ec/mod.py
+++ b/pyecsca/ec/mod.py
@@ -51,6 +51,7 @@ def check(func):
 
     return method
 
+
 @public
 class RandomModAction(Action):
     """A random sampling from Z_n."""