-
Notifications
You must be signed in to change notification settings - Fork 0
/
SceneGraph.py
116 lines (90 loc) · 3.93 KB
/
SceneGraph.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
from PygameUtils import rotate_around
# A transformer is simply a function that takes in an entity and its parent,
# and modifies the child in some way
def NoopTransformer(*args, **kwargs):
"""Doesn't do anything"""
pass
def OldStyleTransformer(entity, parent_entity):
entity.unrotated_position[0] = entity.position[0]
entity.unrotated_position[1] = entity.position[1]
parent_position = parent_entity.absolute_position
# Translate the entity based on parent's position
entity.unrotated_position[0] += parent_position[0]
entity.unrotated_position[1] += parent_position[1]
# Rotate around parent's absolute_position
entity.absolute_position = rotate_around(parent_entity.cos_radians,
parent_entity.sin_radians,
entity.unrotated_position,
parent_entity.absolute_position,
parent_entity.heading)
class SceneGraph(object):
def __init__(self):
super().__init__()
self.nodes = []
self.entity_node_map = {}
def add(self, entity):
"""Add to top level"""
self.nodes.append(GraphNode(entity))
def add_to_parent(self, entity, parent, transformer=NoopTransformer):
"""
Recursively search scene graph for node containing 'parent'
Add entity to parent
"""
parent = self.find_node_by_entity(parent, self.nodes)
if not parent:
raise Exception("Node for parent entity '{}' couldn't be found!".format(entity))
parent.add_child_entity(entity, transformer)
def update(self):
for node in self.nodes:
node.update()
def find_node_by_entity(self, entity, root_nodes=None):
"""Recursively search scene graph for node containing entity"""
if root_nodes is None:
root_nodes = self.nodes
if len(root_nodes) == 0:
return None
for node in root_nodes:
if node.entity == entity:
return node
child_search = self.find_node_by_entity(entity, node.children)
if child_search:
return child_search
return None
def remove(self, entity):
"""Find the node and remove it from it's parent if it has one"""
node = self.find_node_by_entity(entity, self.nodes)
if not node:
raise Exception("Node for entity '{}' couldn't be found!".format(entity))
if node.parent:
node.parent.remove_child_entity(entity)
else:
self.nodes.remove(node)
def __contains__(self, entity):
return self.find_node_by_entity(entity, self.nodes) is not None
class GraphNode(object):
def __init__(self, entity, parent=None, transformer=NoopTransformer):
super().__init__()
self.entity = entity
self.parent = parent
self.transformer = transformer
self.children = []
def update(self):
"""
Transform this node, then update the children nodes recursively
"""
if self.parent:
self.transformer(self.entity, self.parent.entity)
for child in self.children:
child.update()
def add_child_entity(self, child_entity, transformer=NoopTransformer):
new_node = GraphNode(child_entity, parent=self, transformer=transformer)
self.children.append(new_node)
def find_child_by_entity(self, entity):
"""Find the first child GraphNode which contains entity, or return None"""
return next((child for child in self.children if child.entity == entity), None)
def remove_child_entity(self, child_entity):
found_node = self.find_child_by_entity(child_entity)
if found_node:
self.children.remove(found_node)
else:
raise Exception("Child node for entity '{}' couldn't be found!".format(child_entity))