rigify/rig_lists.py

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