rigify/rig_lists.py

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