Extend parsing of Scientific Notation to JSON capabilites

Author: maximilian-tech

src/datamodel_code_generator/__init__.py

@@ -30,7 +30,7 @@
     PythonVersionMin,
 )
 from datamodel_code_generator.parser import DefaultPutDict, LiteralType
-from datamodel_code_generator.util import SafeLoader
+from datamodel_code_generator.util import SafeLoader, get_safeloader
 
 MIN_VERSION: Final[int] = 9
 MAX_VERSION: Final[int] = 13
@@ -46,9 +46,14 @@
 
 DEFAULT_BASE_CLASS: str = "pydantic.BaseModel"
 
+g_extend_yaml_scientifc_notation: bool = False
+
+def get_safeloader_wrapper(safeloader: type[SafeLoader]) -> type[SafeLoader]:
+    return get_safeloader(safeloader, extend_yaml_scientifc_notation=g_extend_yaml_scientifc_notation)
+
 
 def load_yaml(stream: str | TextIO) -> Any:
-    return yaml.load(stream, Loader=SafeLoader)  # noqa: S506
+    return yaml.load(stream, Loader=get_safeloader_wrapper(SafeLoader))  # noqa: S506
 
 
 def load_yaml_from_path(path: Path, encoding: str) -> Any:
@@ -262,6 +267,7 @@ def generate(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
     use_title_as_name: bool = False,
     use_operation_id_as_name: bool = False,
     use_unique_items_as_set: bool = False,
+    extend_yaml_scientifc_notation: bool = False,
     http_headers: Sequence[tuple[str, str]] | None = None,
     http_ignore_tls: bool = False,
     use_annotated: bool = False,
@@ -288,6 +294,8 @@ def generate(  # noqa: PLR0912, PLR0913, PLR0914, PLR0915
     no_alias: bool = False,
     formatters: list[Formatter] = DEFAULT_FORMATTERS,
 ) -> None:
+    global g_extend_yaml_scientifc_notation
+    g_extend_yaml_scientifc_notation = extend_yaml_scientifc_notation
     remote_text_cache: DefaultPutDict[str, str] = DefaultPutDict()
     if isinstance(input_, str):
         input_text: str | None = input_

src/datamodel_code_generator/__main__.py

@@ -290,6 +290,7 @@ def validate_root(cls, values: Any) -> Any:  # noqa: N805
     use_title_as_name: bool = False
     use_operation_id_as_name: bool = False
     use_unique_items_as_set: bool = False
+    extend_yaml_scientifc_notation: bool = False
     http_headers: Optional[Sequence[tuple[str, str]]] = None  # noqa: UP045
     http_ignore_tls: bool = False
     use_annotated: bool = False
@@ -498,6 +499,7 @@ def main(args: Sequence[str] | None = None) -> Exit:  # noqa: PLR0911, PLR0912,
             use_title_as_name=config.use_title_as_name,
             use_operation_id_as_name=config.use_operation_id_as_name,
             use_unique_items_as_set=config.use_unique_items_as_set,
+            extend_yaml_scientifc_notation=config.extend_yaml_scientifc_notation,
             http_headers=config.http_headers,
             http_ignore_tls=config.http_ignore_tls,
             use_annotated=config.use_annotated,

src/datamodel_code_generator/arguments.py

@@ -286,6 +286,12 @@ def start_section(self, heading: str | None) -> None:
     action="store_true",
     default=None,
 )
+typing_options.add_argument(
+    "--extend-yaml-scientifc-notation",
+    help="Parse scientific floating point as specified by json instead of yaml.",
+    action="store_true",
+    default=None,
+)
 
 # ======================================================================================
 # Customization options for generated model fields

src/datamodel_code_generator/util.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import copy
+import re
 from typing import TYPE_CHECKING, Any, Callable, TypeVar
 
 import pydantic
@@ -34,14 +35,30 @@ def load_toml(path: Path) -> dict[str, Any]:
         with path.open("rb") as f:
             return load_tomllib(f)
 
+def get_safeloader(safeloader: type[SafeLoader], *, extend_yaml_scientifc_notation: bool) -> type[SafeLoader]:
+    safeloadertemp = copy.deepcopy(safeloader)
+    safeloadertemp.yaml_constructors = copy.deepcopy(safeloader.yaml_constructors)
+    safeloadertemp.add_constructor(
+        "tag:yaml.org,2002:timestamp",
+        safeloadertemp.yaml_constructors["tag:yaml.org,2002:str"],
+    )
+    if extend_yaml_scientifc_notation is True:
+        safeloadertemp.add_implicit_resolver(
+            "tag:yaml.org,2002:float",
+            re.compile(
+                r"""^(?:
+            [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
+            |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
+            |\\.[0-9_]+(?:[eE][-+][0-9]+)?
+            |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
+            |[-+]?\\.(?:inf|Inf|INF)
+            |\\.(?:nan|NaN|NAN))$""",
+                re.VERBOSE,
+            ),
+            list("-+0123456789."),
+        )
+    return safeloadertemp
 
-SafeLoaderTemp = copy.deepcopy(SafeLoader)
-SafeLoaderTemp.yaml_constructors = copy.deepcopy(SafeLoader.yaml_constructors)
-SafeLoaderTemp.add_constructor(
-    "tag:yaml.org,2002:timestamp",
-    SafeLoaderTemp.yaml_constructors["tag:yaml.org,2002:str"],
-)
-SafeLoader = SafeLoaderTemp
 
 Model = TypeVar("Model", bound=_BaseModel)