From c79c61f79b95c3a3f2adccdfb23173a2263a18c2 Mon Sep 17 00:00:00 2001
From: Auric Vente <mprops@protonmail.com>
Date: Thu, 25 Jul 2024 21:23:38 -0600
Subject: [PATCH] Mods

---
 cromulant/args.py    |  2 ++
 cromulant/argspec.py |  8 +++++++-
 cromulant/storage.py | 34 +++++++++++++++++++++++++++-------
 3 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/cromulant/args.py b/cromulant/args.py
index 58fb192..7fc9f21 100644
--- a/cromulant/args.py
+++ b/cromulant/args.py
@@ -10,6 +10,7 @@ from .argspec import ArgSpec
 
 class Args:
     names: Path | None = None
+    ants: Path | None = None
     images: bool = True
     header: bool = True
     footer: bool = True
@@ -51,6 +52,7 @@ class Args:
 
         paths = [
             "names",
+            "ants",
         ]
 
         for p_item in paths:
diff --git a/cromulant/argspec.py b/cromulant/argspec.py
index bc07351..8bbfbb0 100644
--- a/cromulant/argspec.py
+++ b/cromulant/argspec.py
@@ -73,7 +73,13 @@ class ArgSpec:
         ArgSpec.add_argument(
             "names",
             type=str,
-            info="Path to a JSON file with a list of names. The game will use these names instead of the default ones",
+            info="Path to a JSON file with a list of names. Use these instead of the default ones",
+        )
+
+        ArgSpec.add_argument(
+            "ants",
+            type=str,
+            info="Path to a JSON file with ants data. Use this instead of the default one",
         )
 
         ArgSpec.add_argument(
diff --git a/cromulant/storage.py b/cromulant/storage.py
index 6333d58..4266892 100644
--- a/cromulant/storage.py
+++ b/cromulant/storage.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import json
 from typing import TYPE_CHECKING, Any
+from pathlib import Path
 
 from .config import Config
 
@@ -13,10 +14,32 @@ from .utils import Utils
 
 
 class Storage:
+    @staticmethod
+    def get_names_path() -> Path:
+        path = Config.names_json
+
+        if Args.names:
+            if Args.names.exists():
+                path = Args.names
+
+        return path
+
+    @staticmethod
+    def get_ants_path() -> Path:
+        path = Config.ants_json
+
+        if Args.ants:
+            if Args.ants.exists():
+                path = Args.ants
+
+        return path
+
     @staticmethod
     def get_ants() -> Any:
         try:
-            with Config.ants_json.open() as file:
+            path = Storage.get_ants_path()
+
+            with path.open() as file:
                 return json.load(file)
         except Exception as e:
             Utils.print(str(e))
@@ -25,17 +48,14 @@ class Storage:
     @staticmethod
     def save_ants(ants: list[Ant]) -> None:
         objs = [ant.to_dict() for ant in ants]
+        path = Storage.get_ants_path()
 
-        with Config.ants_json.open("w") as file:
+        with path.open("w") as file:
             json.dump(objs, file)
 
     @staticmethod
     def get_names() -> Any:
-        path = Config.names_json
-
-        if Args.names:
-            if Args.names.exists():
-                path = Args.names
+        path = Storage.get_names_path()
 
         with path.open() as file:
             return json.load(file)