space_view3d_math_vis/utils.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 import bpy
   4
   5
   6 def console_namespace():
   7     import console_python
   8     for window in bpy.context.window_manager.windows:
   9         for area in window.screen.areas:
  10             if area.type == 'CONSOLE':
  11                 for region in area.regions:
  12                     if region.type == 'WINDOW':
  13                         console = console_python.get_console(hash(region))
  14                         if console:
  15                             return console[0].locals
  16     return {}
  17
  18
  19 def is_display_list(listvar):
  20     from mathutils import Vector
  21
  22     for var in listvar:
  23         if not isinstance(var, Vector):
  24             return False
  25     return True
  26
  27
  28 class VarStates:
  29
  30     @staticmethod
  31     def store_states():
  32         # Store the display states, called upon unregister the Add-on
  33         # This is useful when you press F8 to reload the Addons.
  34         # Then this function preserves the display states of the
  35         # console variables.
  36         state_props = bpy.context.window_manager.MathVisStatePropList
  37         variables = get_math_data()
  38
  39         for index, state_prop in reversed(list(enumerate(state_props))):
  40             if state_prop.name not in variables:
  41                 # Variable has been removed from console
  42                 state_props.remove(index)
  43
  44         for key, ktype in variables.items():
  45             if key and key not in state_props:
  46                 prop = state_props.add()
  47                 prop.name = key
  48                 prop.ktype = ktype.__name__
  49                 prop.state = [True, False]
  50
  51     @staticmethod
  52     def get_index(key):
  53         index = bpy.context.window_manager.MathVisStatePropList.find(key)
  54         return index
  55
  56     @staticmethod
  57     def delete(key):
  58         state_props = bpy.context.window_manager.MathVisStatePropList
  59         index = state_props.find(key)
  60         if index != -1:
  61             state_props.remove(index)
  62
  63     @staticmethod
  64     def toggle_display_state(key):
  65         state_props = bpy.context.window_manager.MathVisStatePropList
  66         if key in state_props:
  67             state_props[key].state[0] = not state_props[key].state[0]
  68         else:
  69             print("Odd: Can not find key %s in MathVisStateProps" % (key))
  70
  71     @staticmethod
  72     def toggle_lock_state(key):
  73         state_props = bpy.context.window_manager.MathVisStatePropList
  74         if key in state_props:
  75             state_props[key].state[1] = not state_props[key].state[1]
  76         else:
  77             print("Odd: Can not find key %s in MathVisStateProps" % (key))
  78
  79
  80 def get_math_data():
  81     from mathutils import Matrix, Vector, Quaternion, Euler
  82
  83     locals = console_namespace()
  84     if not locals:
  85         return {}
  86
  87     variables = {}
  88     for key, var in locals.items():
  89         if len(key) == 0 or key[0] == "_":
  90             continue
  91
  92         type_var = type(var)
  93
  94         # Rules out sets/dicts.
  95         # It's also possible the length check below is slow
  96         # for data with underlying linked-list structure.
  97         if not hasattr(type_var, "__getitem__"):
  98             continue
  99
 100         # Don't do a truth test on the data because this causes an error with some
 101         # array types, see T66107.
 102         len_fn = getattr(type_var, "__len__", None)
 103         if len_fn is None:
 104             continue
 105         if len_fn(var) == 0:
 106             continue
 107
 108         if isinstance(var, (Matrix, Vector, Quaternion, Euler)) or \
 109            isinstance(var, (tuple, list)) and is_display_list(var):
 110
 111             variables[key] = type_var
 112
 113     return variables
 114
 115
 116 def cleanup_math_data():
 117
 118     locals = console_namespace()
 119     if not locals:
 120         return
 121
 122     variables = get_math_data()
 123
 124     for key in variables.keys():
 125         index = VarStates.get_index(key)
 126         if index == -1:
 127             continue
 128
 129         state_prop = bpy.context.window_manager.MathVisStatePropList.get(key)
 130         if state_prop.state[1]:
 131             continue
 132
 133         del locals[key]
 134         bpy.context.window_manager.MathVisStatePropList.remove(index)
 135
 136
 137 def console_math_data():
 138     from mathutils import Matrix, Vector, Quaternion, Euler
 139
 140     data_matrix = {}
 141     data_quat = {}
 142     data_euler = {}
 143     data_vector = {}
 144     data_vector_array = {}
 145
 146     for key, var in console_namespace().items():
 147         if key[0] == "_":
 148             continue
 149
 150         state_prop = bpy.context.window_manager.MathVisStatePropList.get(key)
 151         if state_prop:
 152             disp, lock = state_prop.state
 153             if not disp:
 154                 continue
 155
 156         if isinstance(var, Matrix):
 157             if len(var.col) != 4 or len(var.row) != 4:
 158                 if len(var.col) == len(var.row):
 159                     var = var.to_4x4()
 160                 else:  # todo, support 4x3 matrix
 161                     continue
 162             data_matrix[key] = var
 163         elif isinstance(var, Vector):
 164             if len(var) < 3:
 165                 var = var.to_3d()
 166             data_vector[key] = var
 167         elif isinstance(var, Quaternion):
 168             data_quat[key] = var
 169         elif isinstance(var, Euler):
 170             data_euler[key] = var
 171         elif type(var) in {list, tuple} and is_display_list(var):
 172             data_vector_array[key] = var
 173
 174     return data_matrix, data_quat, data_euler, data_vector, data_vector_array