space_view3d_math_vis/utils.py

   1 # ##### BEGIN GPL LICENSE BLOCK #####
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 # ##### END GPL LICENSE BLOCK #####
  18
  19 # <pep8 compliant>
  20
  21 import bpy
  22
  23
  24 def console_namespace():
  25     import console_python
  26     get_consoles = console_python.get_console
  27     consoles = getattr(get_consoles, "consoles", None)
  28     if consoles:
  29         for console, stdout, stderr in get_consoles.consoles.values():
  30             return console.locals
  31     return {}
  32
  33
  34 def is_display_list(listvar):
  35     from mathutils import Vector
  36
  37     for var in listvar:
  38         if type(var) is not Vector:
  39             return False
  40     return True
  41
  42
  43 class VarStates:
  44
  45     @staticmethod
  46     def store_states():
  47         # Store the display states, called upon unregister the Add-on
  48         # This is useful when you press F8 to reload the Addons.
  49         # Then this function preserves the display states of the
  50         # console variables.
  51         state_props = bpy.context.window_manager.MathVisStatePropList
  52         variables = get_math_data()
  53         for key, ktype in variables.items():
  54             if key and key not in state_props:
  55                 prop = state_props.add()
  56                 prop.name = key
  57                 prop.ktype = ktype.__name__
  58                 prop.state = [True, False]
  59
  60     @staticmethod
  61     def get_index(key):
  62         index = bpy.context.window_manager.MathVisStatePropList.find(key)
  63         return index
  64
  65     @staticmethod
  66     def delete(key):
  67         state_props = bpy.context.window_manager.MathVisStatePropList
  68         index = state_props.find(key)
  69         if index != -1:
  70             state_props.remove(index)
  71
  72     @staticmethod
  73     def toggle_display_state(key):
  74         state_props = bpy.context.window_manager.MathVisStatePropList
  75         if key in state_props:
  76             state_props[key].state[0] = not state_props[key].state[0]
  77         else:
  78             print("Odd: Can not find key %s in MathVisStateProps" % (key))
  79
  80     @staticmethod
  81     def toggle_lock_state(key):
  82         state_props = bpy.context.window_manager.MathVisStatePropList
  83         if key in state_props:
  84             state_props[key].state[1] = not state_props[key].state[1]
  85         else:
  86             print("Odd: Can not find key %s in MathVisStateProps" % (key))
  87
  88
  89 def get_math_data():
  90     from mathutils import Matrix, Vector, Quaternion, Euler
  91
  92     locals = console_namespace()
  93     if not locals:
  94         return {}
  95
  96     variables = {}
  97     for key, var in locals.items():
  98         if key[0] == "_" or not var:
  99             continue
 100         if type(var) in {Matrix, Vector, Quaternion, Euler} or \
 101            type(var) in {tuple, list} and is_display_list(var):
 102
 103             variables[key] = type(var)
 104
 105     return variables
 106
 107
 108 def cleanup_math_data():
 109
 110     locals = console_namespace()
 111     if not locals:
 112         return
 113
 114     variables = get_math_data()
 115
 116     for key in variables.keys():
 117         index = VarStates.get_index(key)
 118         if index == -1:
 119             continue
 120
 121         state_prop = bpy.context.window_manager.MathVisStatePropList.get(key)
 122         if state_prop.state[1]:
 123             continue
 124
 125         del locals[key]
 126         bpy.context.window_manager.MathVisStatePropList.remove(index)
 127
 128
 129 def console_math_data():
 130     from mathutils import Matrix, Vector, Quaternion, Euler
 131
 132     data_matrix = {}
 133     data_quat = {}
 134     data_euler = {}
 135     data_vector = {}
 136     data_vector_array = {}
 137
 138     for key, var in console_namespace().items():
 139         if key[0] == "_":
 140             continue
 141
 142         state_prop = bpy.context.window_manager.MathVisStatePropList.get(key)
 143         if state_prop:
 144             disp, lock = state_prop.state
 145             if not disp:
 146                 continue
 147
 148         var_type = type(var)
 149
 150         if var_type is Matrix:
 151             if len(var.col) != 4 or len(var.row) != 4:
 152                 if len(var.col) == len(var.row):
 153                     var = var.to_4x4()
 154                 else:  # todo, support 4x3 matrix
 155                     continue
 156             data_matrix[key] = var
 157         elif var_type is Vector:
 158             if len(var) < 3:
 159                 var = var.to_3d()
 160             data_vector[key] = var
 161         elif var_type is Quaternion:
 162             data_quat[key] = var
 163         elif var_type is Euler:
 164             data_euler[key] = var
 165         elif var_type in {list, tuple} and is_display_list(var):
 166             data_vector_array[key] = var
 167
 168     return data_matrix, data_quat, data_euler, data_vector, data_vector_array