space_view3d_math_vis/utils.py

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