rigify/rig_lists.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 import os
   4 import traceback
   5 import importlib
   6 import typing
   7
   8 from typing import Optional, Iterable
   9
  10 from .utils.rig import RIG_DIR
  11
  12 from . import feature_set_list
  13
  14
  15 def get_rigs(base_dir: str, base_path: list[str], *,
  16              path: Iterable[str] = (),
  17              feature_set=feature_set_list.DEFAULT_NAME):
  18     """ Recursively searches for rig types, and returns a list.
  19
  20     Args:
  21         base_dir:      root directory
  22         base_path:     base dir where rigs are stored
  23         path:          rig path inside the base dir
  24         feature_set:   feature set that is being loaded
  25     """
  26
  27     rig_table = {}
  28     impl_rigs = {}
  29
  30     dir_path = os.path.join(base_dir, *path)
  31
  32     try:
  33         files = os.listdir(dir_path)
  34     except FileNotFoundError:
  35         files = []
  36
  37     files.sort()
  38
  39     for f in files:
  40         is_dir = os.path.isdir(os.path.join(dir_path, f))  # Whether the file is a directory
  41
  42         # Stop cases
  43         if f[0] in [".", "_"]:
  44             continue
  45         if f.count(".") >= 2 or (is_dir and "." in f):
  46             print("Warning: %r, filename contains a '.', skipping" % os.path.join(*base_path, *path, f))
  47             continue
  48
  49         if is_dir:
  50             # Check for sub-rigs
  51             sub_rigs, sub_impls = get_rigs(base_dir, base_path, path=[*path, f], feature_set=feature_set)
  52             rig_table.update(sub_rigs)
  53             impl_rigs.update(sub_impls)
  54         elif f.endswith(".py"):
  55             # Check straight-up python files
  56             sub_path = [*path, f[:-3]]
  57             key = '.'.join(sub_path)
  58             # Don't reload rig modules - it breaks isinstance
  59             rig_module = importlib.import_module('.'.join(base_path + sub_path))
  60             if hasattr(rig_module, "Rig"):
  61                 rig_table[key] = {"module": rig_module,
  62                                   "feature_set": feature_set}
  63             if hasattr(rig_module, 'IMPLEMENTATION') and rig_module.IMPLEMENTATION:
  64                 impl_rigs[key] = rig_module
  65
  66     return rig_table, impl_rigs
  67
  68
  69 # Public variables
  70 rigs = {}
  71 implementation_rigs = {}
  72
  73
  74 def get_rig_class(name: str) -> Optional[typing.Type]:
  75     try:
  76         return rigs[name]["module"].Rig
  77     except (KeyError, AttributeError):
  78         return None
  79
  80
  81 def get_internal_rigs():
  82     global rigs, implementation_rigs
  83
  84     base_rigify_dir = os.path.dirname(__file__)
  85     base_rigify_path = __name__.split('.')[:-1]
  86
  87     rigs, implementation_rigs = get_rigs(os.path.join(base_rigify_dir, RIG_DIR),
  88                                          [*base_rigify_path, RIG_DIR])
  89
  90
  91 def get_external_rigs(set_list):
  92     # Clear and fill rigify rigs and implementation rigs public variables
  93     for rig in list(rigs.keys()):
  94         if rigs[rig]["feature_set"] != feature_set_list.DEFAULT_NAME:
  95             rigs.pop(rig)
  96             if rig in implementation_rigs:
  97                 implementation_rigs.pop(rig)
  98
  99     # Get external rigs
 100     for feature_set in set_list:
 101         # noinspection PyBroadException
 102         try:
 103             base_dir, base_path = feature_set_list.get_dir_path(feature_set, RIG_DIR)
 104
 105             external_rigs, external_impl_rigs = get_rigs(base_dir, base_path, feature_set=feature_set)
 106         except Exception:
 107             print(f"Rigify Error: Could not load feature set '{feature_set}' rigs: exception occurred.\n")
 108             traceback.print_exc()
 109             print("")
 110             continue
 111
 112         rigs.update(external_rigs)
 113         implementation_rigs.update(external_impl_rigs)