From a0cbf613b9d3ffe6633a590caf2700fe8f8972d2 Mon Sep 17 00:00:00 2001 From: Max Wiklund Date: Wed, 7 Sep 2022 23:31:41 -0700 Subject: [PATCH] Add support for custom maya command modules. --- README.md | 13 +++++++++---- mayaff/__init__.py | 11 ++++++++--- mayaff/config.py | 5 ++++- mayaff/pyparser.py | 18 +++++++++--------- tests/test_pyparser.py | 21 +++++++++++++++++++++ 5 files changed, 51 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index d1c6215..92c19df 100644 --- a/README.md +++ b/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 diff --git a/mayaff/__init__.py b/mayaff/__init__.py index 3dfc417..6019005 100644 --- a/mayaff/__init__.py +++ b/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,6 +67,11 @@ 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() @@ -75,7 +79,8 @@ def set_up_argparser() -> argparse.Namespace: 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: diff --git a/mayaff/config.py b/mayaff/config.py index fbcb958..5d4fb0e 100644 --- a/mayaff/config.py +++ b/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. diff --git a/mayaff/pyparser.py b/mayaff/pyparser.py index 31cb5b2..29f657f 100644 --- a/mayaff/pyparser.py +++ b/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 diff --git a/tests/test_pyparser.py b/tests/test_pyparser.py index f35c0a3..db7fe1f 100644 --- a/tests/test_pyparser.py +++ b/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))