From 7b6e6e2845619805cbfa534e54e7ba4f7370d14b Mon Sep 17 00:00:00 2001
From: Kristian Hartikainen <kristian.hartikainen@gmail.com>
Date: Tue, 15 Oct 2024 10:46:47 -0400
Subject: [PATCH] Make mtl parsing more robust using regex

Now supports e.g. whitespaces in filename
---
 obj2mjcf/cli.py   | 22 +++++++++++++++-------
 tests/test_cli.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/obj2mjcf/cli.py b/obj2mjcf/cli.py
index 80b433e..084b0b9 100644
--- a/obj2mjcf/cli.py
+++ b/obj2mjcf/cli.py
@@ -1,5 +1,6 @@
 """A CLI for processing composite Wavefront OBJ files for use in MuJoCo."""
 
+from collections.abc import Iterable
 import logging
 import os
 import re
@@ -110,6 +111,15 @@ def decompose_convex(filename: Path, work_dir: Path, coacd_args: CoacdArgs) -> b
     return True
 
 
+def parse_mtl_name(lines: Iterable[str]) -> Optional[str]:
+    mtl_regex = re.compile(r"^mtllib\s+(.+?\.mtl)(?:\s*#.*)?\s*\n?$")
+    for line in lines:
+        match = mtl_regex.match(line)
+        if match is not None:
+            name = match.group(1)
+            return name
+    return None
+
 def process_obj(filename: Path, args: Args) -> None:
     # Create a directory with the same name as the OBJ file. The processed submeshes
     # and materials will be stored there.
@@ -133,13 +143,11 @@ def process_obj(filename: Path, args: Args) -> None:
 
     # Check if the OBJ files references an MTL file.
     # TODO(kevin): Should we support multiple MTL files?
-    process_mtl = False
-    with open(filename, "r") as f:
-        for line in f.readlines():
-            if line.startswith("mtllib"):  # Deals with commented out lines.
-                process_mtl = True
-                name = line.split()[1]
-                break
+    with filename.open("r") as f:
+        name = parse_mtl_name(f.readlines())
+
+    process_mtl = name is not None
+
 
     sub_mtls: List[List[str]] = []
     mtls: List[Material] = []
diff --git a/tests/test_cli.py b/tests/test_cli.py
index c256484..cbe50c7 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -1,6 +1,9 @@
+from itertools import product
 import pathlib
 import subprocess
 
+from obj2mjcf import cli
+
 # Path to the directory containing this file.
 _THIS_DIR = pathlib.Path(__file__).parent.absolute()
 
@@ -18,3 +21,34 @@ def test_runs_without_error() -> None:
         ]
     )
     assert retcode == 0
+
+
+def test_parse_mtl_line() -> None:
+    file_names = [
+        "A B.mtl",
+        "A_B.mtl",
+        "AbCd.mtl",
+        "a.mtl",
+        "a b.mtl",
+        "a-b.mtl",
+        "a_b.mtl",
+    ]
+    comments = [
+        "",
+        "# comment",
+        " # comment",
+        " # comment #",
+    ]
+    for file_name, comment in product(file_names, comments):
+        line = f"mtllib {file_name}{comment}\n"
+        result = cli.parse_mtl_name(iter([line]))
+        assert result == file_name, result
+
+    for line in (
+            "a",
+            "a # b",
+            "mtllib a.what",
+            "a.mtl",
+    ):
+        result = cli.parse_mtl_name(iter([line]))
+        assert result is None, result