Fix issue with shortcut recompression

montepy/input_parser/input_reader.py

@@ -3,15 +3,17 @@
 from montepy.constants import DEFAULT_VERSION
 
 
-def read_input(input_file, mcnp_version=DEFAULT_VERSION, replace=True):
+def read_input(destination, mcnp_version=DEFAULT_VERSION, replace=True):
     """
     Reads the specified MCNP Input file.
 
     The MCNP version must be a three component tuple e.g., (6, 2, 0) and (5, 1, 60).
 
+    .. note::
+        if a stream is provided. It will not be closed by this function.
 
-    :param input_file: the path to the input file to read.
-    :type input_file: str
+    :param destination: the path to the input file to read, or a readable stream.
+    :type destination: io.TextIOBase, str, os.PathLike
     :param mcnp_version: The version of MCNP that the input is intended for.
     :type mcnp_version: tuple
     :returns: The MCNP_Problem instance representing this file.
@@ -24,7 +26,7 @@ def read_input(input_file, mcnp_version=DEFAULT_VERSION, replace=True):
     :raises BrokenObjectLinkError: If a reference is made to an object that is not in the input file.
     :raises UnknownElement: If an isotope is specified for an unknown element.
     """
-    problem = mcnp_problem.MCNP_Problem(input_file)
+    problem = mcnp_problem.MCNP_Problem(destination)
     problem.mcnp_version = mcnp_version
     problem.parse_input(replace=replace)
     return problem

montepy/input_parser/syntax_node.py

@@ -1834,7 +1834,7 @@ def _can_consume_node(self, node, direction, last_edge_shortcut=False):
         :type node: ValueNode
         :param direction: the direct to go in. Must be in {-1, 1}
         :type direction: int
-        :param last_edge_shortcut: Whether or the previous node in the list was
+        :param last_edge_shortcut: Whether or not the previous node in the list was
             part of a different shortcut
         :type last_edge_shortcut: bool
         :returns: true it can be consumed.
@@ -1967,7 +1967,24 @@ def _format_jump(self):
         return f"{num_jumps.format()}{j}"
 
     def _format_repeat(self, leading_node=None):
-        if leading_node is not None:
+        def can_use_last_node(node):
+            """Last node can be used if
+            - it's a basic ValueNode
+            - it's a different shortcut that concludes itself well.
+            - it's also a repeat, with the same edge values.
+            """
+            return isinstance(node, ValueNode) or (
+                isinstance(node, type(self))
+                and (
+                    leading_node._type != Shortcuts.REPEAT
+                    or (
+                        leading_node._type == Shortcuts.REPEAT
+                        and math.isclose(self.nodes[0].value, node.nodes[-1].value)
+                    )
+                )
+            )
+
+        if can_use_last_node(leading_node):
             first_val = ""
             num_extra = 0
         else:

montepy/mcnp_problem.py

@@ -23,12 +23,18 @@ class MCNP_Problem:
     """
     A class to represent an entire MCNP problem in a semantic way.
 
-    :param file_name: the path to the file that will be read.
-    :type file_name: str
+    .. note::
+        If a stream is provided. It will not be closed by this function.
+
+    :param destination: the path to the input file to read, or a readable stream.
+    :type destination: io.TextIOBase, str, os.PathLike
     """
 
-    def __init__(self, file_name):
-        self._input_file = MCNP_InputFile(file_name)
+    def __init__(self, destination):
+        if hasattr(destination, "read") and callable(getattr(destination, "read")):
+            self._input_file = MCNP_InputFile.from_open_stream(destination)
+        elif isinstance(destination, (str, os.PathLike)):
+            self._input_file = MCNP_InputFile(destination)
         self._title = None
         self._message = None
         self._print_in_data_block = CellDataPrintController()

tests/inputs/test_importance.imcnp

@@ -41,7 +41,7 @@ MT3 lwtr.23t
 C execution
 ksrc 0 0 0
 imp:n,p 1 1 1 0 3
-imp:e   0 0 0 1 2
+imp:e   0 2r 1 r
 kcode 100000 1.000 50 1050
 phys:p j 1 2j 1
 mode n p e

tests/test_integration.py

@@ -1,6 +1,7 @@
 # Copyright 2024, Battelle Energy Alliance, LLC All Rights Reserved.
 import copy
 import io
+from pathlib import Path
 
 import pytest
 import os
@@ -17,6 +18,8 @@
 from montepy.particle import Particle
 import numpy as np
 
+from tests import constants
+
 
 @pytest.fixture(scope="module")
 def simple_problem():
@@ -481,7 +484,7 @@ def test_importance_format_unmutated(importance_problem):
     print(output)
     assert len(output) == 2
     assert "imp:n,p 1 1 1 0 3" == output[0]
-    assert "imp:e   0 0 0 1 2" == output[1]
+    assert "imp:e   0 2r 1 r" == output[1]
 
 
 def test_importance_format_mutated(importance_problem):
@@ -1066,3 +1069,31 @@ def test_alternate_encoding():
     montepy.read_input(
         os.path.join("tests", "inputs", "bad_encoding.imcnp"), replace=True
     )
+
+
+@pytest.mark.parametrize(
+    "file",
+    {"tests/inputs/test_importance.imcnp"},
+    # set((Path("tests") / "inputs").iterdir())
+    # - {
+    #    Path("tests") / "inputs" / p
+    #    for p in constants.BAD_INPUTS | constants.IGNORE_FILES
+    # },
+)
+def test_read_write_cycle(file):
+    print(f"Testing against {file} *********************")
+    problem = montepy.read_input(file)
+    print(problem.cells._importance._real_tree[montepy.particle.Particle.ELECTRON])
+    print()
+    fh = io.StringIO()
+    problem.write_problem(fh)
+    fh.seek(0)
+    # test valid syntax
+    new_problem = montepy.MCNP_Problem("foo")
+    # verify lines are similar
+    fh.seek(0)
+    lines = [line.rstrip() for line in fh]
+    [print(line) for line in lines]
+    with open(file, "r") as gold_fh:
+        for gold_line, new_line in zip(gold_fh, lines):
+            assert new_line == gold_line.rstrip().expandtabs(8)