pymcnp/cli: updated io

Author: BitterB0NG0

src/demo/main.py

@@ -1,25 +1,19 @@
-import sys
-
-import docopt
-
-
-DEMO_DOCOPT = """
+"""
 Usage:
     demo <inp>
 """
 
+import docopt
+
 
 def main():
-    args = docopt.docopt(DEMO_DOCOPT)
+    args = docopt.docopt(__doc__)
 
     try:
         open(args['<inp>'])
-        print(f"[DEMO] Args: {' '.join(sys.argv)}")
     except FileNotFoundError:
-        print('[DEMO] Input file not found.')
         exit(1)
 
     path = ''.join(args['<inp>'].split('.')[:-1]) + '.outp'
     with open(path, 'x') as file:
-        print(path)
         file.write('MCNP Output! :)')

src/pymcnp/cli/_errors.py

@@ -11,14 +11,15 @@ class CliCode(enum.Enum):
     Represents ``CliError`` error codes.
 
     Notes:
-        10X: Value, ``Run``.
+        10X: System.
+        11X: Value, ``Run``.
     """
 
-    INVALID_INP = 100
-    INVALID_PATH = 101
-    INVALID_PREHOOK = 102
-    INVALID_POSTHOOK = 103
-    INVALID_COMMAND = 104
+    INVALID_INP = 110
+    INVALID_PATH = 111
+    INVALID_PREHOOK = 112
+    INVALID_POSTHOOK = 113
+    INVALID_COMMAND = 114
 
 
 class CliError(Exception):
@@ -73,6 +74,6 @@ def __str__(self) -> str:
                 hint += ''
 
         if hint:
-            return f'\n\033[31;4;1mCliError[{self.code.name}]\033[0m: {head}\n|\n| {repr(self.info)}\n|\n| \033[35;4mHint\033[0m: {hint}\n|'
+            return f'[red][bold]McnpError[{self.code.name}]:[/][/] {head}\n\n``{self.info[:]}``\n\n[magenta][bold]Hint:[/bold][/magenta] {hint}'
         else:
-            return f'\n\033[31;4;1mCliError[{self.code.name}]\033[0m: {head}\n|\n| {repr(self.info)}\n|'
+            return f'[red][bold]McnpError[{self.code.name}]:[/][/] {head}\n\n``{self.info[:]}``'

src/pymcnp/cli/_io.py

@@ -5,17 +5,31 @@
 import datetime
 
 import rich
+import rich.panel
 
 
 def get_timestamp() -> str:
     return datetime.datetime.today().strftime('%Y-%m-%d--%H-%M-%S')
 
 
 def error(msg):
-    rich.print(msg)
+    rich.print(rich.panel.Panel(msg))
     exit(1)
 
 
-def warning(msg):
+def print(msg):
     rich.print(rich.panel.Panel(msg))
-    return
+
+
+def done():
+    rich.print(rich.panel.Panel('[deep_sky_blue1][bold]Info:[/][/] Done!'))
+
+
+def warning():
+    rich.print(
+        rich.panel.Panel(
+            '[bold][gold3]Warning:[/][/] PyMCNP is just getting started! '
+            'Please, double check the output to make sure everything works. '
+            'If you find errors, please, report them on [link=https://github.com/FSIBT/PyMCNP/issues]GitHub[/link].'
+        )
+    )

src/pymcnp/cli/check.py

@@ -1,20 +1,18 @@
 """
 Usage:
-    pymcnp check [options] <input_file>
+    pymcnp check [options] <file>
 
 Options:
     -f --fix                  Reformat input file.
-    -o file --output=file     Write to file instead of overwriting the input file.
 """
 
-from pathlib import Path
-import sys
+import pathlib
+import difflib
 
 from docopt import docopt
-from rich import print
-from rich.panel import Panel
 
-import pymcnp
+from . import _io
+from .. import files
 
 
 def main() -> None:
@@ -24,42 +22,27 @@ def main() -> None:
     ``pymcnp check`` checks INP file syntax.
     """
 
-    args = docopt(__doc__)
-
-    input_file = Path(args['<input_file>'])
-
-    print(
-        Panel(
-            '[orange3]Warning[/] PyMCNP is just getting started.'
-            'Please double check the output to make sure everyhing is working as expected'
-            'If you find an error, please report it at https://github.com/FSIBT/PyMCNP/issues'
-        )
-    )
+    _io.warning()
 
-    if not input_file.is_file():
-        print(f'[red]ERROR[/] Input file {input_file} does not exists.')
-        sys.exit(1)
+    # Processing CLI arguments.
+    args = docopt(__doc__)
+    file = pathlib.Path(args['<file>'])
 
+    # Reading INP file.
     try:
-        data = pymcnp.read_input(input_file)
-    except:  # noqa
-        print('[red]ERROR[/] Cannot read input file.')
-        print(
-            '     We would appreciate if you can report this '
-            '(and if possible share the input file) at https://github.com/FSIBT/PyMCNP/issues'
-            'Thanks! -- your PyMCNP team'
-        )
-        sys.exit(2)
-
+        inp = files.inp.Inp.from_mcnp_file(file)
+    except files.utils.errors.McnpError as err:
+        _io.error(err.__str__())
+    except FileNotFoundError:
+        _io.error(f'[red][bold]IoError:[/][/] ``{file}`` not found.')
+
+    # Running ``diff``!
+    with open(file, 'r') as file:
+        diff = difflib.unified_diff(file.readlines(), inp.to_mcnp().split('\n'))
+        _io.print('\n'.join(line.rstrip('\n') for line in diff))
+
+    # Handling ``--fix``.
     if args['--fix']:
-        output = args['--output']
-        if output is None:
-            print(data.to_mcnp())
-        else:
-            output = Path(output)
-            if output.is_file():
-                print('[orange3]Warning[/] Overwriting existing file.')
-            data.to_mcnp_file(output)
-            print(':thumbs_up: Rewrote input file.')
-    else:
-        print(f':thumbs_up: Input file {input_file} can be parsed.')
+        inp.to_mcnp_file(file)
+
+    _io.done()

src/pymcnp/cli/convert.py

@@ -20,9 +20,9 @@
 
 from docopt import docopt
 from rich import print
-from rich.panel import Panel
 
 import pymcnp
+from . import _io
 
 
 def main() -> None:
@@ -32,19 +32,13 @@ def main() -> None:
     ``pymcnp convert`` converts MCNP output files to pandas dataframes.
     """
 
+    _io.warning()
+
     args = docopt(__doc__)
 
     N = args['--number']
     file_in = Path(args['<input>'])
 
-    print(
-        Panel(
-            '[orange3]Warning[/] PyMCNP is just getting started.'
-            'Please double check the output to make sure everyhing is working as expected'
-            'If you find an error, please report it at https://github.com/FSIBT/PyMCNP/issues'
-        )
-    )
-
     if not file_in.is_file():
         print(f'[red]Error[/] Cannot access file {file_in}')
         sys.exit(1)
@@ -69,3 +63,5 @@ def main() -> None:
                 print(f'Saving tally {N} as parquet to ', file_out)
             case _:
                 print('Currently cannot handle this file type')
+
+    _io.done()

src/pymcnp/cli/hooks/__init__.py

@@ -0,0 +1,9 @@
+from . import nps
+from . import clean
+from . import empty
+
+__all__ = [
+    'nps',
+    'clean',
+    'empty',
+]

src/pymcnp/cli/hooks/clean.py

@@ -1,13 +1,11 @@
+import shutil
 import pathlib
 
 
 def main(path):
     path = pathlib.Path(path)
 
     for file_or_dir in path.rglob('*'):
-        if file_or_dir.is_file():
-            file_or_dir.unlink()
-        else:
-            file_or_dir.rmdir()
+        shutil.rmtree(file_or_dir)
 
     path.rmdir()

src/pymcnp/cli/hooks/default.py src/pymcnp/cli/hooks/empty.pysrc/pymcnp/cli/hooks/default.py renamed to src/pymcnp/cli/hooks/empty.py

@@ -13,9 +13,9 @@
 from docopt import docopt
 from matplotlib import pyplot as plt
 from rich import print
-from rich.panel import Panel
 
 import pymcnp
+from . import _io
 
 
 def main() -> None:
@@ -25,19 +25,13 @@ def main() -> None:
     ``pymcnp plot`` plots MCNP output data.
     """
 
+    _io.warning()
+
     args = docopt(__doc__)
 
     N = args['--number']
     file_in = Path(args['<input>'])
 
-    print(
-        Panel(
-            '[orange3]Warning[/] PyMCNP is just getting started.'
-            'Please double check the output to make sure everyhing is working as expected'
-            'If you find an error, please report it at https://github.com/FSIBT/PyMCNP/issues'
-        )
-    )
-
     if not file_in.is_file():
         print(f'[red]Error[/] Cannot access file {file_in}')
         sys.exit(1)
@@ -78,3 +72,5 @@ def main() -> None:
         plt.savefig(output_filename)
     else:
         plt.show()
+
+    _io.done()

src/pymcnp/cli/plot.py

@@ -16,9 +16,9 @@
 from docopt import docopt
 
 from . import _io
+from . import hooks
 from . import _errors
 from .. import files
-from .hooks import default
 
 
 EXECUTABLE = """
@@ -57,15 +57,17 @@ class Run:
         path_subdirectories: List of paths to run directories.
     """
 
-    PREHOOK_PASS = default
-    POSTHOOK_PASS = default
+    PREHOOK_EMPTY = hooks.empty
+    POSTHOOK_EMPTY = hooks.empty
+    PREHOOK_NPS = hooks.nps
+    POSTHOOK_CLEAN = hooks.clean
 
     def __init__(
         self,
         inps: list[files.inp.Inp],
         path: pathlib.Path,
-        prehook: ModuleType = PREHOOK_PASS,
-        posthook: ModuleType = POSTHOOK_PASS,
+        prehook: ModuleType = PREHOOK_EMPTY,
+        posthook: ModuleType = POSTHOOK_EMPTY,
         command: str = 'mcnp6',
     ):
         """
@@ -127,7 +129,7 @@ def __init__(
 
     def run(self, hosts: list[str] = []):
         """
-        Runs MCNP INP files.
+        Runs MCNP INP pymcnp.
 
         Parameters:
             hosts: List of hostnames on which to execute.
@@ -168,25 +170,35 @@ def main() -> None:
     """
     Executes the ``pymcnp run`` command.
 
-    ``pymcnp run`` runs INP files.
+    ``pymcnp run`` runs INP pymcnp.
     """
 
-    args = docopt(__doc__)
+    _io.warning()
 
+    # Processing CLI arguments.
+    args = docopt(__doc__)
     inps = args['<file>']
     hosts = args['--hosts'] if args['--hosts'] else []
     command = args['--command'] if args['--command'] else 'mcnp6'
 
+    # Reading INP files.
     try:
         inps = [files.inp.Inp.from_mcnp_file(inp) for inp in inps]
     except files.utils.errors.McnpError as err:
         _io.error(err.__str__())
+    except FileNotFoundError:
+        _io.error(f'[red][bold]IoError:[/][/] {inps} File not found.')
 
-    path = pathlib.Path(os.getcwd())
-
+    # Running!
     try:
-        run = Run(inps, path, command=command)
+        Run(
+            inps,
+            pathlib.Path(os.getcwd()),
+            prehook=Run.PREHOOK_EMPTY,
+            posthook=Run.POSTHOOK_CLEAN,
+            command=command,
+        ).run(hosts)
     except _errors.CliError as err:
         _io.error(err.__str__())
 
-    run.run(hosts)
+    _io.done()