From e86bdfcbd883c89946bbc7699ef643d6f99dc0b8 Mon Sep 17 00:00:00 2001
From: Simon Schoonjans <simon.schoonjans@waylay.io>
Date: Tue, 26 Mar 2024 03:02:15 +0100
Subject: [PATCH] feat: support option to prefer default and/or examples values
 (#110)

---
 jsf/parser.py                            | 15 ++++++---
 jsf/schema_types/base.py                 |  7 +++++
 jsf/tests/data/object-with-examples.json | 21 +++++++++++++
 jsf/tests/test_default_fake.py           | 40 ++++++++++++++++++++++++
 4 files changed, 78 insertions(+), 5 deletions(-)
 create mode 100644 jsf/tests/data/object-with-examples.json

diff --git a/jsf/parser.py b/jsf/parser.py
index 0c83613..3174429 100644
--- a/jsf/parser.py
+++ b/jsf/parser.py
@@ -339,16 +339,21 @@ def _parse(self, schema: Dict[str, Any]) -> AllTypes:
     def context(self):
         return {**self.base_context, "state": deepcopy(self.base_state)}
 
-    def generate(self, n: Optional[int] = None) -> Any:
+    def generate(
+        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.
 
-        If n is provided, it returns a list of n objects. If n is 1 then
-        it returns a single object.
+        Args:
+            n (int, optional): If n is provided, it returns a list of n objects. If n is 1 then it returns a single object.
+            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}
         if n is None or n == 1:
-            return self.root.generate(context=self.context)
-        return [self.root.generate(context=self.context) for _ in range(n)]
+            return self.root.generate(context=context)
+        return [self.root.generate(context=context) for _ in range(n)]
 
     def pydantic(self):
         """Generates a fake object from the provided schema and provides the
diff --git a/jsf/schema_types/base.py b/jsf/schema_types/base.py
index 85e5976..3edd443 100644
--- a/jsf/schema_types/base.py
+++ b/jsf/schema_types/base.py
@@ -56,6 +56,13 @@ def generate(self, context: Dict[str, Any]) -> Any:
             return None
         if self.provider is not None:
             return eval(self.provider, context)()
+
+        if context.get("use_defaults", False) and self.default:
+            return self.default
+
+        if context.get("use_examples", False) and self.examples:
+            return random.choice(self.examples)
+
         raise ProviderNotSetException()
 
     def model(self, context: Dict[str, Any]) -> Optional[Tuple[Type, Field]]:
diff --git a/jsf/tests/data/object-with-examples.json b/jsf/tests/data/object-with-examples.json
new file mode 100644
index 0000000..3d37732
--- /dev/null
+++ b/jsf/tests/data/object-with-examples.json
@@ -0,0 +1,21 @@
+{
+    "title": "Pet",
+    "type": "object",
+    "properties": {
+      "name": {
+        "type": "string",
+        "examples": ["Chop", "Luna", "Thanos"]
+      },
+      "species": {
+        "type": "string",
+        "examples": ["Dog", "Cat", "Rabbit"]
+      },
+      "breed": {
+        "type": "string",
+        "default": "Mixed Breed",
+        "examples": ["Labrador Retriever", "Siamese", "Golden Retriever"]
+      }
+    },
+    "required": ["name", "species"]
+  }
+  
\ No newline at end of file
diff --git a/jsf/tests/test_default_fake.py b/jsf/tests/test_default_fake.py
index 3a9c69c..56b42e6 100644
--- a/jsf/tests/test_default_fake.py
+++ b/jsf/tests/test_default_fake.py
@@ -504,3 +504,43 @@ def test_fake_empty(TestData):
     with open(TestData / "empty.json") as file:
         schema = json.load(file)
     [JSF(schema).generate() for _ in range(10)]  # Just validating no errors
+
+
+def test_use_defaults(TestData):
+    with open(TestData / "object-with-examples.json") as file:
+        schema = json.load(file)
+    p = JSF(schema)
+
+    fake_data = [p.generate(use_defaults=True) for _ in range(10)]
+    for d in fake_data:
+        assert isinstance(d, dict)
+        breed = d.get("breed")
+        assert breed is None or breed == "Mixed Breed"
+
+
+def test_use_examples(TestData):
+    with open(TestData / "object-with-examples.json") as file:
+        schema = json.load(file)
+    p = JSF(schema)
+
+    fake_data = [p.generate(use_examples=True) for _ in range(10)]
+    for d in fake_data:
+        assert isinstance(d, dict)
+        assert d["species"] in ["Dog", "Cat", "Rabbit"]
+        assert d["name"] in ["Chop", "Luna", "Thanos"]
+        breed = d.get("breed")
+        assert breed is None or breed in ["Labrador Retriever", "Siamese", "Golden Retriever"]
+
+
+def test_use_defaults_and_examples(TestData):
+    with open(TestData / "object-with-examples.json") as file:
+        schema = json.load(file)
+    p = JSF(schema)
+
+    fake_data = [p.generate(use_defaults=True, use_examples=True) for _ in range(10)]
+    for d in fake_data:
+        assert isinstance(d, dict)
+        assert d["species"] in ["Dog", "Cat", "Rabbit"]
+        assert d["name"] in ["Chop", "Luna", "Thanos"]
+        breed = d.get("breed")
+        assert breed is None or breed == "Mixed Breed"