Add support for custom maya command modules.

README.md

@@ -58,10 +58,10 @@ mayaff only supports python 3.7+
 
 ## Usage
 Options:
-```bash
+```
 usage: mayaff [-h] [-v] [-t {2018}] [--check] [--diff] [--quiet] [--exclude EXCLUDE]
-           [--exclude-files EXCLUDE_FILES [EXCLUDE_FILES ...]] [--single-thread]
-           source [source ...]
+              [--exclude-files EXCLUDE_FILES [EXCLUDE_FILES ...]] [--modules MODULES] [--single-thread]
+              source [source ...]
 
 Command line tool to find and replace short maya flags.
 
@@ -79,13 +79,18 @@ optional arguments:
   --exclude EXCLUDE     A regular expression for file names to exclude.
   --exclude-files EXCLUDE_FILES [EXCLUDE_FILES ...]
                         Exclude files. Separate files with space.
+  --modules MODULES     Maya modules to use for import. Example: --modules 'maya:cmds,pymel:core'
   --single-thread       Only execute mayaff on single thread.
 ```
 
-```bash
+```
 mayaff . --exclude-files package.py
 ```
 
+Some companies have custom wrappers around `cmds`. If you want `mayaff` to find your cmds module somewhere else use `--modules`.
+Example: `mayaff . --modules "maya:cmds,custom.module:cmds`
+
+
 ## Known limitations
 
 ```python

mayaff/__init__.py

@@ -24,7 +24,7 @@
 from mayaff.config import MayaArgsConfig
 
 
-__version__ = "0.1.0-beta.1"
+__version__ = "0.1.0-beta.2"
 _DESCRIPTION = "Command line tool to find and replace short maya flags."
 
 
@@ -34,7 +34,6 @@ def set_up_argparser() -> argparse.Namespace:
         f.replace(".json", "") for f in pkg_resources.resource_listdir("mayaff", "maya_configs") if f.endswith(".json")
     ]
 
-    # config_options = []
     parser = argparse.ArgumentParser(description=_DESCRIPTION)
     parser.add_argument("-v", "--version", action="version", version="%(prog)s {}".format(__version__))
     parser.add_argument("source", nargs="+", help="Directory or files you want to format.")
@@ -68,14 +67,20 @@ def set_up_argparser() -> argparse.Namespace:
         help="A regular expression for file names to exclude.",
     )
     parser.add_argument("--exclude-files", nargs="+", default=[], help="Exclude files. Separate files with space.")
+    parser.add_argument(
+        "--modules",
+        default="maya:cmds,pymel:core",
+        help="Maya modules to use for import. Examples: --modules 'maya:cmds,pymel:core'",
+    )
     parser.add_argument("--single-thread", action="store_true", default=False, help="Only execute mayaff on single thread.")
     return parser.parse_args()
 
 
 def _main() -> int:
     """Run command line app with return code."""
     args = set_up_argparser()
-    config = MayaArgsConfig(args.target_version)
+    modules = [tuple(m.split(":")) for m in args.modules.split(",")]
+    config = MayaArgsConfig(args.target_version, modules)
     try:
         exclued_re = re.compile(args.exclude)
     except re.error:

mayaff/config.py

@@ -14,20 +14,23 @@
 
 import json
 import pkgutil
+from typing import List, Tuple
 
 
 class MayaArgsConfig(object):
     """Class to manage maya flags configuration."""
 
-    def __init__(self, config_version: str = "2018"):
+    def __init__(self, config_version: str = "2018", modules: List[Tuple[str, str]] = tuple(("maya", "cmds"))):
         """Construct class and load config.
 
         Args:
             config_version: Configuration name.
+            modules: List of maya modules to look for maya commands.
 
         """
         super().__init__()
         self._command_data = json.loads(pkgutil.get_data("mayaff", f"maya_configs/{config_version}.json"))
+        self.modules = modules
 
     def get_flags(self, command_name: str) -> dict:
         """Try to get command flags from command name.

mayaff/pyparser.py

@@ -80,10 +80,11 @@ def peek(self) -> tokenize.TokenInfo:
 class MayaImportVisitor(ast.NodeVisitor):
     """Ast traversal class to find cmds import from maya."""
 
-    def __init__(self):
+    def __init__(self, modules):
         """Construct class and do nothing."""
         super().__init__()
         self.maya_imports = []
+        self.modules = modules
 
     def visit_Import(self, node: ast.Import) -> None:
         """Code to check if `maya.cmds` is imported.
@@ -93,8 +94,9 @@ def visit_Import(self, node: ast.Import) -> None:
 
         """
         for _import in node.names:
-            if _import.name == "maya.cmds":
-                self.maya_imports.append(_import.asname or _import.name)
+            for imp in self.modules:
+                if _import.name == ".".join(imp):
+                    self.maya_imports.append(_import.asname or _import.name)
 
     def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
         """Code to check if `maya.cmds` is imported.
@@ -103,12 +105,10 @@ def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
             node: Node to check if maya is imported.
 
         """
-        if not node.module == "maya":
-            return
-
         for _import in node.names:
-            if _import.name == "cmds":
-                self.maya_imports.append(_import.asname or _import.name)
+            for module, imp in self.modules:
+                if node.module == module and _import.name == imp:
+                    self.maya_imports.append(_import.asname or _import.name)
 
 
 class MayaFlagsParser(object):
@@ -315,6 +315,6 @@ def _parse_maya_imports(self, source_code: str, file_name: str) -> None:
 
         """
         tree = ast.parse(source_code, file_name)
-        import_visitor = MayaImportVisitor()
+        import_visitor = MayaImportVisitor(self._config.modules)
         import_visitor.visit(tree)
         self._found_maya_modules = import_visitor.maya_imports

tests/test_pyparser.py

@@ -203,3 +203,24 @@ def test_parse_python2_raise_exception(self):
         )
         with self.assertRaises(SyntaxError) as context:
             self.assertRaises(SyntaxError, mayaff_api.format_string(source_code, self.config_cls))
+
+    def test_custom_modules(self):
+        source_code = textwrap.dedent(
+            """
+            from ABC.maya2 import abc
+
+            abc.textureWindow(source, ra=True)
+            """
+        )
+
+        expected_result = textwrap.dedent(
+            """
+            from ABC.maya2 import abc
+
+            abc.textureWindow(source, removeAllImages=True)
+            """
+        )
+
+        _config = config.MayaArgsConfig(modules=[("ABC.maya2", "abc")])
+
+        self.assertEqual(expected_result, mayaff_api.format_string(source_code, config=_config))