From dc844501db73c49752756a046a3e20bd2afca99e Mon Sep 17 00:00:00 2001
From: Marc Harnos <mharnos@gmail.com>
Date: Sun, 1 Jan 2023 22:56:50 +0100
Subject: [PATCH] feat: enable underscore integer separation for decimal values
 (ints)

---
 docs/index.md                   | 30 +++++++++++++++++++++++++++---
 src/yamlfix/adapters.py         | 18 ++++++++++++++++++
 src/yamlfix/model.py            |  1 +
 tests/unit/test_adapter_yaml.py | 22 ++++++++++++++++++++++
 4 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/docs/index.md b/docs/index.md
index ce2f5b6..2c2bfaf 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -315,7 +315,7 @@ none_value5: null
 Default: `quote_basic_values: bool = False`<br>
 Environment variable override:
 ```bash
-export YAMLFIX_quote_basic_values="false"
+export YAMLFIX_QUOTE_BASIC_VALUES="false"
 ```
 
 Per default `ruyaml` will quote only values where it is necessary to explicitly define the type of a value. This is the case for numbers and boolean values for example. If your `yaml`-file contains a value of type string that would look like a number, then this value needs to be quoted.
@@ -341,7 +341,7 @@ stringList: ["abc", "123"]
 Default: `quote_keys_and_basic_values: bool = False`<br>
 Environment variable override:
 ```bash
-export YAMLFIX_quote_keys_and_basic_values="false"
+export YAMLFIX_QUOTE_KEYS_AND_BASIC_VALUES="false"
 ```
 
 Similar to the [quote basic values](#quote-basic-values) configuration option, this option, in addition to the values themselves, quotes the keys as well. For example:
@@ -363,7 +363,7 @@ list: [item, item]
 Default: `quote_representation: str = "'"`<br>
 Environment variable override:
 ```bash
-export YAMLFIX_quote_representation="'"
+export YAMLFIX_QUOTE_REPRESENTATION="'"
 ```
 
 Configure which quotation string is used for quoting values. For example:
@@ -378,6 +378,30 @@ key: 'value'
 key: "value"
 ```
 
+### Underscore Integer
+
+Default: `underscore_integer: bool = False`<br>
+Environment variable override:
+```bash
+export YAMLFIX_UNDERSCORE_INTEGER="true"
+```
+
+Configure whether to add underscores to big integer values or not. For example:
+
+```yaml
+---
+small_integer: 10
+big_integer: 1000
+```
+
+would be converted to
+
+```yaml
+---
+small_integer: 10
+big_integer: 1_000
+```
+
 # References
 
 As most open sourced programs, `yamlfix` is standing on the shoulders of giants,
diff --git a/src/yamlfix/adapters.py b/src/yamlfix/adapters.py
index 3b292e7..7a666f4 100644
--- a/src/yamlfix/adapters.py
+++ b/src/yamlfix/adapters.py
@@ -8,6 +8,7 @@
 from ruyaml.main import YAML
 from ruyaml.nodes import MappingNode, Node, ScalarNode, SequenceNode
 from ruyaml.representer import RoundTripRepresenter
+from ruyaml.scalarint import ScalarInt
 from ruyaml.tokens import CommentToken
 
 from yamlfix.model import YamlfixConfig
@@ -123,6 +124,21 @@ def represent_str(self, data: Any) -> ScalarNode:  # noqa: ANN401
             "tag:yaml.org,2002:str", data, self.config.quote_representation
         )
 
+    def represent_int(self, data: Any) -> Any:  # noqa: ANN401
+        """Configure the SafeRepresenter integer representation to use underscore\
+            separation for big integers."""
+        if not self.config.underscore_integer:
+            return super().represent_int(data)
+
+        return self.represent_scalar("tag:yaml.org,2002:int", f"{data:_d}")
+
+    def represent_scalar_int(self, data: Any) -> Any:  # noqa: ANN401
+        """Configure the RoundTripRepresenter integer representation to use underscore\
+            separation for big integers."""
+        # provide the internal `_underscore` flag, that is used in insert_underscore
+        data._underscore = self.config.underscore_integer  # noqa: W0212
+        return super().represent_scalar_int(data)
+
     def represent_mapping(
         self, tag: Any, mapping: Any, flow_style: Optional[Any] = None  # noqa: ANN401
     ) -> MappingNode:
@@ -314,6 +330,8 @@ def _apply_simple_value_quotations(self, value_node: Node) -> None:
 
 YamlfixRepresenter.add_representer(type(None), YamlfixRepresenter.represent_none)
 YamlfixRepresenter.add_representer(str, YamlfixRepresenter.represent_str)
+YamlfixRepresenter.add_representer(int, YamlfixRepresenter.represent_int)
+YamlfixRepresenter.add_representer(ScalarInt, YamlfixRepresenter.represent_scalar_int)
 
 
 class SourceCodeFixer:
diff --git a/src/yamlfix/model.py b/src/yamlfix/model.py
index e969cac..589289e 100644
--- a/src/yamlfix/model.py
+++ b/src/yamlfix/model.py
@@ -22,3 +22,4 @@ class YamlfixConfig(ConfigSchema):
     quote_basic_values: bool = False
     quote_keys_and_basic_values: bool = False
     quote_representation: str = "'"
+    underscore_integer: bool = False
diff --git a/tests/unit/test_adapter_yaml.py b/tests/unit/test_adapter_yaml.py
index 4ff9759..2d7e9dc 100644
--- a/tests/unit/test_adapter_yaml.py
+++ b/tests/unit/test_adapter_yaml.py
@@ -582,3 +582,25 @@ def test_empty_list_inline_comment_indentation(self) -> None:
         result = fix_code(source, config)
 
         assert result == source
+
+    def test_underscore_for_decimal_integers(self) -> None:
+        """Make underscore for decimal integers configurable."""
+        source = dedent(
+            """\
+            small_ints: [1, 10, 100]
+            big_ints: [1000, 10000, 1000000, 100000000000000000000]
+            """
+        )
+        fixed_source = dedent(
+            """\
+            ---
+            small_ints: [1, 10, 100]
+            big_ints: [1_000, 10_000, 1_000_000, 100_000_000_000_000_000_000]
+            """  # noqa: E501
+        )
+        config = YamlfixConfig()
+        config.underscore_integer = True
+
+        result = fix_code(source, config)
+
+        assert result == fixed_source