DMI web export functions (#7)

Co-authored-by: Metal Gear Sloth <[email protected]>

rsi/__init__.py

@@ -1,4 +1,19 @@
 from .rsi import Rsi
+from .splitters import (
+    HyphenSplitter,
+    NumberSplitter,
+    RsiSplitter,
+    SimpleSplitter,
+    UnderscoreSplitter,
+)
 from .state import State
 
-__all__ = [ "Rsi", "State" ]
+__all__ = [
+    "Rsi",
+    "State" ,
+    "HyphenSplitter",
+    "NumberSplitter",
+    "RsiSplitter",
+    "SimpleSplitter",
+    "UnderscoreSplitter"
+]

rsi/__main__.py

@@ -1,7 +1,15 @@
 import argparse
 from pathlib import Path
-from typing import Optional
-from rsi import Rsi
+from typing import Optional, Union
+from io import BytesIO
+
+from rsi import (
+    Rsi,
+    HyphenSplitter,
+    NumberSplitter,
+    SimpleSplitter,
+    UnderscoreSplitter,
+)
 
 
 def main() -> int:
@@ -12,10 +20,12 @@ def main() -> int:
     subparser = parser.add_subparsers(dest="command")
 
     _from_dmi = subparser.add_parser("from_dmi", help="Will create an RSI from a BYOND DMI file.")
-    _from_dmi.add_argument("input", help="The DMI file to read from.", type=Path)
+    _from_dmi.add_argument("input", help="The DMI file to read from.")
     _from_dmi.add_argument("output", help="The RSI to output to.", type=Path)
     _from_dmi.add_argument("-c", "--copyright", help="Specifies the copyright of the new RSI file.")
+    _from_dmi.add_argument("-i", "--indents", help="Indents to use for the meta.json file.", type=int)
     _from_dmi.add_argument("-l", "--license", help="Specifies the license of the new RSI file.")
+    _from_dmi.add_argument("-s", "--splitter", help="Splitter to use for the rsi if applicable.")
 
     _new_rsi = subparser.add_parser("new", help="Will create a new RSI at the provided directory.")
     _new_rsi.add_argument("rsi", help="The location of the new RSI. Must not exist yet.", type=Path)
@@ -27,7 +37,7 @@ def main() -> int:
     args = parser.parse_args()
 
     if args.command == "from_dmi":
-        from_dmi(args.input, args.output, args.license, args.copyright)
+        from_dmi(args.input, args.output, args.license, args.copyright, args.indents, args.splitter)
         return 0
 
     if args.command == "new":
@@ -37,13 +47,29 @@ def main() -> int:
     return 1
 
 
-def from_dmi(inputf: Path, output: Path, new_license: Optional[str], new_copyright: Optional[str]) -> None:
+def from_dmi(inputf: Union[BytesIO, Path, str],
+             output: Path,
+             new_license: Optional[str],
+             new_copyright: Optional[str],
+             indents: Optional[int] = None,
+             splitter: Optional[str] = "",) -> None:
     rsi = Rsi.from_dmi(inputf)
     if new_license:
         rsi.license = new_license
     if new_copyright:
         rsi.copyright = new_copyright
-    rsi.write(output)
+
+    splitter_class = {
+        "hyphen": HyphenSplitter,
+        "number": NumberSplitter,
+        "simple": SimpleSplitter,
+        "underscore": UnderscoreSplitter,
+    }.get(splitter, None)  # type: ignore
+
+    if splitter_class is None:
+        rsi.write(output)
+    else:
+        splitter_class(rsi).split_to(output, indents)
 
 
 def new_rsi(loc: Path,

rsi/rsi.py

@@ -2,6 +2,8 @@
 import math
 from pathlib import Path
 from typing import Dict, Tuple, Union, cast, TextIO, Any, List, Type, Optional
+from urllib.request import urlopen
+from io import BytesIO
 from PIL import Image
 from .direction import Direction
 from .state import State
@@ -155,20 +157,42 @@ def open(cls, path: Union[str, Path]) -> "Rsi":
         return rsi
 
     @classmethod
-    def from_dmi(cls, path: Union[str, Path]) -> "Rsi":
+    def from_dmi(cls, path: Union[BytesIO, str, Path]) -> "Rsi":
         try:
             from byond.DMI import DMI
         except ImportError:
             raise ImportError("Unable to import byondtoolsv3.")
 
+        # We'll set the copyright to the url if applicable
+        url_copyright = None
+
         if isinstance(path, Path):
             path = str(path)
+        elif isinstance(path, BytesIO):
+            pass
+        else:
+            # Try URL / filepath
+            try:
+                with urlopen(path) as response:
+                    buffer = BytesIO()
+                    buffer.write(response.read())
+                    url_copyright = path
+
+                path = buffer
+            except ValueError:
+                path = Path(path)  # type: ignore
+                if not path.exists():
+                    raise FileNotFoundError
+
+                path = str(path)
 
         # N3X15, if you are reading this:
         # You are awful at API design.
         dmi = DMI(path)
         dmi.loadAll()
         rsi = Rsi((dmi.icon_width, dmi.icon_height))
+        if url_copyright is not None:
+            rsi.copyright = url_copyright
 
         for dmstate in dmi.states.values():
             rsstate = rsi.new_state(dmstate.dirs, dmstate.name)  # type: State

rsi/splitters.py

@@ -0,0 +1,124 @@
+"""
+Used for converting a single .rsi file into mutiple smaller ones
+"""
+from rsi import Rsi
+from typing import Dict, List, Optional
+from abc import ABC
+from pathlib import Path
+import re
+
+
+class RsiSplitter(ABC):
+    """
+    This encapsulates all of the data needed to split a Rsi into several smaller Rsi somewhat cleanly
+    """
+
+    def __init__(self, rsi: Rsi) -> None:
+        self.rsi = rsi
+        # Store the name to use for each rsi group
+        self.names: Dict[Rsi, str] = {}
+
+    def _split(self) -> List[Rsi]:
+        raise NotImplementedError
+
+    def split_to(self, path: Path, indents: Optional[int] = None) -> None:
+        for rsi in self._split():
+            rsi.license = self.rsi.license
+            rsi.copyright = self.rsi.copyright
+            rsi_path = path.joinpath(self.names[rsi]).with_suffix(".rsi")
+            rsi.write(rsi_path, indent=indents)
+
+        self.names.clear()
+
+
+class SimpleSplitter(RsiSplitter):
+    """
+    Split each Rsi state into its own Rsi
+    """
+
+    def _split(self) -> List[Rsi]:
+        result = []
+
+        for name, state in self.rsi.states.items():
+            state_rsi = Rsi(self.rsi.size)
+            state_rsi.set_state(state, name)
+            result.append(state_rsi)
+            self.names[state_rsi] = state.name
+
+        return result
+
+
+class HyphenSplitter(RsiSplitter):
+    """
+    Split each rsi based on hyphens where Rsi states with the same prefix are grouped together
+    e.g. ak-20, ak-40, and ak-60 would all be grouped into ak.
+    """
+
+    def _split(self) -> List[Rsi]:
+        groups: Dict[str, Rsi] = {}
+
+        for name, state in self.rsi.states.items():
+            prefix = name.split("-")[0]
+            suffix = name.split("-")[-1]
+            state_rsi = groups.setdefault(prefix, Rsi(self.rsi.size))
+
+            if prefix != suffix:
+                state.name = suffix
+                state_rsi.set_state(state, suffix)
+                self.names[state_rsi] = prefix
+            else:
+                state_rsi.set_state(state, name)
+                self.names[state_rsi] = name
+
+        return list(groups.values())
+
+
+class UnderscoreSplitter(RsiSplitter):
+    """
+    Like the hyphensplitter but for underscores.
+    """
+
+    def _split(self) -> List[Rsi]:
+        groups: Dict[str, Rsi] = {}
+
+        for name, state in self.rsi.states.items():
+            prefix = name.split("_")[0]
+            suffix = name.split("_")[-1]
+            state_rsi = groups.setdefault(prefix, Rsi(self.rsi.size))
+
+            if prefix != suffix:
+                state.name = suffix
+                state_rsi.set_state(state, suffix)
+                self.names[state_rsi] = prefix
+            else:
+                state_rsi.set_state(state, name)
+                self.names[state_rsi] = name
+
+        return list(groups.values())
+
+
+class NumberSplitter(RsiSplitter):
+    """
+    Splits states based on the suffix number
+    e.g. infected, infected0, infected1 are all in the same rsi
+    """
+
+    def _split(self) -> List[Rsi]:
+        groups: Dict[str, Rsi] = {}
+        pattern = re.compile("([^0-9]*)([0-9]*)")
+
+        for name, state in self.rsi.states.items():
+            match = pattern.match(name)
+            prefix = match.group(1)  # type: ignore
+            suffix = match.group(2) if len(match.groups()) > 1 else ""  # type: ignore
+            state_rsi = groups.setdefault(prefix, Rsi(self.rsi.size))
+
+            if prefix != suffix:
+                state.name = suffix
+                state_rsi.set_state(state, suffix)
+                self.names[state_rsi] = prefix
+            else:
+                state_rsi.set_state(state, name)
+                self.names[state_rsi] = name
+
+        return list(groups.values())