Replace base_state counter with a custom picklable one

Author: Romain Bosa

jsf/parser.py

@@ -4,7 +4,6 @@
 from collections import ChainMap
 from copy import deepcopy
 from datetime import datetime
-from itertools import count
 from pathlib import Path
 from types import MappingProxyType
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -68,7 +67,7 @@ def __init__(
         self.root_schema = schema
         self.definitions = {}
         self.base_state = {
-            "__counter__": count(start=1),
+            "__counter__": _PicklableCounter(),
             "__all_json_paths__": [],
             "__depth__": 0,
             **initial_state,
@@ -107,10 +106,16 @@ def from_json(
         """
         with open(path) as f:
             return JSF(
-                json.load(f), context, initial_state, allow_none_optionals, max_recursive_depth
+                json.load(f),
+                context,
+                initial_state,
+                allow_none_optionals,
+                max_recursive_depth,
             )
 
-    def __parse_primitive(self, name: str, path: str, schema: Dict[str, Any]) -> PrimitiveTypes:
+    def __parse_primitive(
+        self, name: str, path: str, schema: Dict[str, Any]
+    ) -> PrimitiveTypes:
         item_type, is_nullable = self.__is_field_nullable(schema)
         cls = Primitives.get(item_type)
         return cls.from_dict(
@@ -125,7 +130,11 @@ def __parse_primitive(self, name: str, path: str, schema: Dict[str, Any]) -> Pri
         )
 
     def __parse_object(
-        self, name: str, path: str, schema: Dict[str, Any], root: Optional[AllTypes] = None
+        self,
+        name: str,
+        path: str,
+        schema: Dict[str, Any],
+        root: Optional[AllTypes] = None,
     ) -> Object:
         _, is_nullable = self.__is_field_nullable(schema)
         model = Object.from_dict(
@@ -142,20 +151,28 @@ def __parse_object(
         props = []
         for _name, definition in schema.get("properties", {}).items():
             props.append(
-                self.__parse_definition(_name, path=f"{path}/{_name}", schema=definition, root=root)
+                self.__parse_definition(
+                    _name, path=f"{path}/{_name}", schema=definition, root=root
+                )
             )
         model.properties = props
         pattern_props = []
         for _name, definition in schema.get("patternProperties", {}).items():
             pattern_props.append(
-                self.__parse_definition(_name, path=f"{path}/{_name}", schema=definition, root=root)
+                self.__parse_definition(
+                    _name, path=f"{path}/{_name}", schema=definition, root=root
+                )
             )
         model.patternProperties = pattern_props
 
         return model
 
     def __parse_array(
-        self, name: str, path: str, schema: Dict[str, Any], root: Optional[AllTypes] = None
+        self,
+        name: str,
+        path: str,
+        schema: Dict[str, Any],
+        root: Optional[AllTypes] = None,
     ) -> Array:
         _, is_nullable = self.__is_field_nullable(schema)
         arr = Array.from_dict(
@@ -169,11 +186,17 @@ def __parse_array(
             }
         )
         root = arr if root is None else root
-        arr.items = self.__parse_definition(name, f"{path}/items", schema["items"], root=root)
+        arr.items = self.__parse_definition(
+            name, f"{path}/items", schema["items"], root=root
+        )
         return arr
 
     def __parse_tuple(
-        self, name: str, path: str, schema: Dict[str, Any], root: Optional[AllTypes] = None
+        self,
+        name: str,
+        path: str,
+        schema: Dict[str, Any],
+        root: Optional[AllTypes] = None,
     ) -> JSFTuple:
         _, is_nullable = self.__is_field_nullable(schema)
         arr = JSFTuple.from_dict(
@@ -190,7 +213,9 @@ def __parse_tuple(
         arr.items = []
         for i, item in enumerate(schema["items"]):
             arr.items.append(
-                self.__parse_definition(name, path=f"{path}/{name}[{i}]", schema=item, root=root)
+                self.__parse_definition(
+                    name, path=f"{path}/{name}[{i}]", schema=item, root=root
+                )
             )
         return arr
 
@@ -207,9 +232,15 @@ def __is_field_nullable(self, schema: Dict[str, Any]) -> Tuple[str, bool]:
         return item_type, False
 
     def __parse_anyOf(
-        self, name: str, path: str, schema: Dict[str, Any], root: Optional[AllTypes] = None
+        self,
+        name: str,
+        path: str,
+        schema: Dict[str, Any],
+        root: Optional[AllTypes] = None,
     ) -> AnyOf:
-        model = AnyOf(name=name, path=path, max_recursive_depth=self.max_recursive_depth, **schema)
+        model = AnyOf(
+            name=name, path=path, max_recursive_depth=self.max_recursive_depth, **schema
+        )
         root = model if root is None else root
         schemas = []
         for d in schema["anyOf"]:
@@ -218,18 +249,32 @@ def __parse_anyOf(
         return model
 
     def __parse_allOf(
-        self, name: str, path: str, schema: Dict[str, Any], root: Optional[AllTypes] = None
+        self,
+        name: str,
+        path: str,
+        schema: Dict[str, Any],
+        root: Optional[AllTypes] = None,
     ) -> AllOf:
         combined_schema = dict(ChainMap(*schema["allOf"]))
-        model = AllOf(name=name, path=path, max_recursive_depth=self.max_recursive_depth, **schema)
+        model = AllOf(
+            name=name, path=path, max_recursive_depth=self.max_recursive_depth, **schema
+        )
         root = model if root is None else root
-        model.combined_schema = self.__parse_definition(name, path, combined_schema, root=root)
+        model.combined_schema = self.__parse_definition(
+            name, path, combined_schema, root=root
+        )
         return model
 
     def __parse_oneOf(
-        self, name: str, path: str, schema: Dict[str, Any], root: Optional[AllTypes] = None
+        self,
+        name: str,
+        path: str,
+        schema: Dict[str, Any],
+        root: Optional[AllTypes] = None,
     ) -> OneOf:
-        model = OneOf(name=name, path=path, max_recursive_depth=self.max_recursive_depth, **schema)
+        model = OneOf(
+            name=name, path=path, max_recursive_depth=self.max_recursive_depth, **schema
+        )
         root = model if root is None else root
         schemas = []
         for d in schema["oneOf"]:
@@ -247,13 +292,20 @@ def __parse_named_definition(self, path: str, def_name: str, root) -> AllTypes:
             definition = schema.get(def_tag, {}).get(def_name)
             if definition is not None:
                 parsed_definition = self.__parse_definition(
-                    def_name, path=f"{path}/#/{def_tag}/{def_name}", schema=definition, root=root
+                    def_name,
+                    path=f"{path}/#/{def_tag}/{def_name}",
+                    schema=definition,
+                    root=root,
                 )
                 self.definitions[f"#/{def_tag}/{def_name}"] = parsed_definition
         return parsed_definition
 
     def __parse_definition(
-        self, name: str, path: str, schema: Dict[str, Any], root: Optional[AllTypes] = None
+        self,
+        name: str,
+        path: str,
+        schema: Dict[str, Any],
+        root: Optional[AllTypes] = None,
     ) -> AllTypes:
         self.base_state["__all_json_paths__"].append(path)
         item_type, is_nullable = self.__is_field_nullable(schema)
@@ -264,7 +316,8 @@ def __parse_definition(
             enum_list = schema["enum"]
             assert len(enum_list) > 0, "Enum List is Empty"
             assert all(
-                isinstance(item, (int, float, str, dict, type(None))) for item in enum_list
+                isinstance(item, (int, float, str, dict, type(None)))
+                for item in enum_list
             ), "Enum Type is not null, int, float, string or dict"
             return JSFEnum.from_dict(
                 {
@@ -286,7 +339,9 @@ def __parse_definition(
             elif item_type == "object" and "oneOf" in schema:
                 return self.__parse_oneOf(name, path, schema, root)
             elif item_type == "array":
-                if (schema.get("contains") is not None) or isinstance(schema.get("items"), dict):
+                if (schema.get("contains") is not None) or isinstance(
+                    schema.get("items"), dict
+                ):
                     return self.__parse_array(name, path, schema, root)
                 if isinstance(schema.get("items"), list) and all(
                     isinstance(x, dict) for x in schema.get("items", [])
@@ -320,7 +375,9 @@ def __parse_definition(
         elif "oneOf" in schema:
             return self.__parse_oneOf(name, path, schema, root)
         elif not any(key in schema for key in ["not", "if", "then", "else"]):
-            return self.__parse_primitive(name, path, {**schema, "type": list(Primitives.keys())})
+            return self.__parse_primitive(
+                name, path, {**schema, "type": list(Primitives.keys())}
+            )
         else:
             raise ValueError(f"Cannot parse schema {repr(schema)}")  # pragma: no cover
 
@@ -340,7 +397,11 @@ def context(self):
         return {**self.base_context, "state": deepcopy(self.base_state)}
 
     def generate(
-        self, n: Optional[int] = None, *, use_defaults: bool = False, use_examples: bool = False
+        self,
+        n: Optional[int] = None,
+        *,
+        use_defaults: bool = False,
+        use_examples: bool = False,
     ) -> Any:
         """Generates a fake object from the provided schema, and returns the
         output.
@@ -350,7 +411,11 @@ def generate(
             use_defaults (bool, optional): prefer the default value as defined in the schema over a randomly generated object. Defaults to False.
             use_examples (bool, optional): prefer an example as defined in the schema over a randomly generated object. This parameter is preceded by the `use_defaults` parameter if set. Defaults to False.
         """
-        context = {**self.context, "use_defaults": use_defaults, "use_examples": use_examples}
+        context = {
+            **self.context,
+            "use_defaults": use_defaults,
+            "use_examples": use_examples,
+        }
         if n is None or n == 1:
             return self.root.generate(context=context)
         return [self.root.generate(context=context) for _ in range(n)]
@@ -371,3 +436,15 @@ def to_json(self, path: Path, **kwargs) -> None:
         output to the given path."""
         with open(path, "w") as f:
             json.dump(self.generate(), f, **kwargs)
+
+
+class _PicklableCounter:
+    def __init__(self):
+        self._counter = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        self._counter += 1
+        return self._counter