/src/examples/Level2/InteractiveFeatureTree/occ_model.py

  1#!/usr/bin/env python
  2
  3##Copyright 2009-2011, Bryan Cole (bryancole.cam@googlemail.com)
  4##
  5##This file is part of pythonOCC.
  6##
  7##pythonOCC is free software: you can redistribute it and/or modify
  8##it under the terms of the GNU Lesser General Public License as published by
  9##the Free Software Foundation, either version 3 of the License, or
 10##(at your option) any later version.
 11##
 12##pythonOCC is distributed in the hope that it will be useful,
 13##but WITHOUT ANY WARRANTY; without even the implied warranty of
 14##MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 15##GNU Lesser General Public License for more details.
 16##
 17##You should have received a copy of the GNU Lesser General Public License
 18##along with pythonOCC.  If not, see <http://www.gnu.org/licenses/>.
 19
 20
 21from enthought.traits.api import (HasTraits, Property, Bool, 
 22                        on_trait_change, cached_property, Instance,
 23                        List, Str, Enum)
 24
 25from enthought.traits.ui.api import View, Item
 26
 27from utils import Tuple, EditorTraits, Int, Float, Range
 28
 29from OCC import TDF, TopoDS, BRepPrimAPI, BRepAlgoAPI, gp, BRepFilletAPI,\
 30        TNaming, TopTools
 31        
 32from OCC.Utils.Topology import Topo
 33
 34###defining an dedicated trait for filter inputs means 
 35###we can track input changes easily
 36Input = Instance(klass="ProcessObject", process_input=True)
 37    
 38
 39class ProcessObject(HasTraits):
 40    """
 41    Base class for all model component objects
 42    """
 43    name = Str
 44    
 45    #
 46    #This flag indicates if the object parameters have changed
 47    #
 48    modified = Bool(True)
 49    
 50    #
 51    #We could link each process-object to a node in an OCAF document
 52    #Not used yet.
 53    #
 54    label = Instance(TDF.TDF_Label)
 55    
 56    #
 57    #Parent TDF_label under which this label will be created
 58    #
 59    parent_label = Instance(TDF.TDF_Label)
 60    
 61    #
 62    #This is the output of the object. The property calls the execute method
 63    #to evaluate the result (which in turn calls up the tree)
 64    #
 65    shape = Property(Instance(TopoDS.TopoDS_Shape))
 66    
 67    #
 68    #Shadow trait which stores the cached shape
 69    #
 70    _shape = Instance(TopoDS.TopoDS_Shape)
 71
 72    #
 73    #A list of all inputs, for the benefit of the TreeEditor
 74    #
 75    _inputs = List
 76      
 77    #
 78    #We hook up listeners to each input to listen to changes in their
 79    #modification trait. Hence, modifications propagate down the tree
 80    #
 81    @on_trait_change("+process_input")
 82    def on_input_change(self, obj, name, vold, vnew):
 83        print "ch", vold, vnew
 84        if vold is not None:
 85            vold.on_trait_change(self.on_modify, 'modified', remove=True)
 86            if vold in self._input_set:
 87                del self._inputs[self._inputs.index(vold)]
 88        
 89        vnew.on_trait_change(self.on_modify, 'modified')
 90        self._inputs.append(vnew)
 91        
 92    def _parent_label_changed(self, old_label, new_label):
 93        ts = TDF.TDF_TagSource()
 94        self.label = ts.NewChild(new_label)
 95        
 96    def on_modify(self, vnew):
 97        if vnew:
 98            self.modified = False
 99            self.modified = True
100        
101    def _get_shape(self):
102        if self.modified:
103            shape = self.execute()
104            self._shape = shape
105            self.modified = False
106            return shape
107        else:
108            return self._shape
109        
110    def execute(self):
111        """return a TopoDS_Shape object"""
112        raise NotImplementedError
113    
114    def update_naming(self, make_shape):
115        """called within the Execute method, to update the Naming
116        Structure. This is key to getting Topological Naming to work"""
117        raise NotImplementedError
118        
119        
120class TopologySource(ProcessObject):
121    def update_naming(self, make_shape):
122        """
123        Create named shapes for the created primitive and all the
124        sub-shapes we want to track. It seems sufficient to 
125        track faces. TNaming_Selector automatically identifies
126        other topology (edges, vertices) based on the faces to
127        which they belong
128        """
129        label = self.label
130        shape = make_shape.Shape()
131        builder = TNaming.TNaming_Builder(label)
132        builder.Generated(shape)
133        for i, face in enumerate(Topo(shape).faces()):
134            f_label = label.FindChild(i+1) #creates a new label if it is not found
135            builder = TNaming.TNaming_Builder(f_label)
136            builder.Generated(face)
137            
138            
139class FilterSource(ProcessObject):
140    def update_naming(self, make_shape):
141        label = self.label
142        shape = make_shape.Shape()
143        
144        input_shape = make_shape.Shape()
145        
146        builder = TNaming.TNaming_Builder(label)
147        builder.Generated(input_shape, shape)
148        
149        #FindChild creates a new label, if one doesn't exist.
150        #Label entry numbers are not necessarily incremental.
151        #They are more like dictionary keys.
152        gen_label = label.FindChild(1)
153        mod_label = label.FindChild(2)
154        del_label = label.FindChild(3)
155        
156        gen_builder = TNaming.TNaming_Builder(gen_label)
157        mod_builder = TNaming.TNaming_Builder(mod_label)
158        del_builder = TNaming.TNaming_Builder(del_label)
159        
160        topo = Topo(input_shape)
161        
162        for face in topo.faces():
163            gen_shapes = make_shape.Generated(face)
164            itr = TopTools.TopTools_ListIteratorOfListOfShape(gen_shapes)
165            while itr.More():
166                this = itr.Value()
167                gen_builder.Generated(face, this)
168                print "generated", face, this
169                itr.Next()
170                        
171        for face in topo.faces():
172            mod_shapes = make_shape.Modified(face)
173            itr = TopTools.TopTools_ListIteratorOfListOfShape(mod_shapes)
174            while itr.More():
175                this = itr.Value()
176                mod_builder.Modified(face, this)
177                print "modified", face, this
178                itr.Next()
179                        
180        for face in topo.faces():
181            if make_shape.IsDeleted(face):
182                del_builder.Delete(face)
183      
184      
185class BlockSource(TopologySource):
186    name = "Block"
187    dims = Tuple(10.0,20.0,30.0, editor_traits={'cols':3})
188    position = Tuple(0.,0.,0., editor_traits={'cols':3})
189    x_axis = Tuple(1.,0.,0., editor_traits={'cols':3})
190    z_axis = Tuple(0.,0.,1., editor_traits={'cols':3})
191    
192    traits_view = View('name',
193                       'dims',
194                       'position',
195                       'x_axis',
196                       'z_axis',
197                       'modified')
198    
199    @on_trait_change("dims, position, x_axis, z_axis")
200    def on_edit(self):
201        self.modified = False
202        self.modified = True
203
204    def execute(self):
205        ax = gp.gp_Ax2(gp.gp_Pnt(*self.position),
206                        gp.gp_Dir(*self.z_axis),
207                        gp.gp_Dir(*self.x_axis))
208        m_box = BRepPrimAPI.BRepPrimAPI_MakeBox(ax, *self.dims)
209        self.update_naming(m_box)
210        return m_box.Shape()
211        
212        
213class SphereSource(TopologySource):
214    name="Sphere"
215    radius = Float(5.0)
216    position = Tuple(0.,0.,0., editor_traits={'cols':3})
217    
218    traits_view = View('name',
219                       'radius',
220                       'position',
221                       'modified')
222    
223    @on_trait_change("radius, position")
224    def on_edit(self):
225        self.modified = False
226        self.modified = True
227        
228    def execute(self):
229        pt = gp.gp_Pnt(*self.position)
230        R = self.radius
231        sph = BRepPrimAPI.BRepPrimAPI_MakeSphere(pt, R)
232        self.update_naming(sph)
233        return sph.Shape()
234        
235        
236class BooleanOpFilter(FilterSource):
237    name = "Boolean Operation"
238    input = Input
239    tool = Input
240    
241    operation = Enum("cut", "fuse", "common")
242    
243    map = {'cut': BRepAlgoAPI.BRepAlgoAPI_Cut,
244           'fuse': BRepAlgoAPI.BRepAlgoAPI_Fuse,
245           'common': BRepAlgoAPI.BRepAlgoAPI_Common}
246    
247    traits_view = View('operation',
248                       'modified')
249    
250    def _operation_changed(self, vnew):
251        self.name = "Boolean Op: %s"%vnew
252        self.modified = False
253        self.modified = True
254        
255    def execute(self):
256        builder = self.map[self.operation]
257        s1 = self.input.shape
258        s2 = self.tool.shape
259        bld = builder(s1, s2)
260        self.update_naming(bld)
261        return bld.Shape()
262        
263    def update_naming(self, make_shape):
264        label = self.label
265        shape = make_shape.Shape()
266        
267        input_shape = make_shape.Shape1()
268        tool_shape = make_shape.Shape2()
269        
270        builder = TNaming.TNaming_Builder(label)
271        builder.Generated(input_shape, shape)
272        builder.Generated(tool_shape, shape)
273        
274        gen_label = label.FindChild(1)
275        mod_label = label.FindChild(2)
276        del_label = label.FindChild(3)
277        
278        gen_builder = TNaming.TNaming_Builder(gen_label)
279        mod_builder = TNaming.TNaming_Builder(mod_label)
280        del_builder = TNaming.TNaming_Builder(del_label)
281        
282        if make_shape.HasGenerated():
283            for in_shape in [input_shape, tool_shape]:
284                for face in Topo(in_shape).faces():
285                    gen_shapes = make_shape.Generated(face)
286                    itr = TopTools.TopTools_ListIteratorOfListOfShape(gen_shapes)
287                    while itr.More():
288                        this = itr.Value()
289                        gen_builder.Generated(face, this)
290                        print "generated", face, this
291                        itr.Next()
292                        
293        if make_shape.HasModified():
294            for face in Topo(input_shape).faces():
295                mod_shapes = make_shape.Modified(face)
296                itr = TopTools.TopTools_ListIteratorOfListOfShape(mod_shapes)
297                while itr.More():
298                    this = itr.Value()
299                    mod_builder.Modify(face, this)
300                    print "modified", face, this
301                    itr.Next()
302                    
303            for face in Topo(tool_shape).faces():
304                mod_shapes = make_shape.Modified2(face)
305                itr = TopTools.TopTools_ListIteratorOfListOfShape(mod_shapes)
306                while itr.More():
307                    this = itr.Value()
308                    mod_builder.Modify(face, this)
309                    print "modified2", face, this
310                    itr.Next()
311                        
312        if make_shape.HasDeleted():
313            for face in Topo(input_shape).faces():
314                if make_shape.IsDeleted(face):
315                    del_builder.Delete(face)
316            for face in Topo(tool_shape).faces():
317                if make_shape.IsDeleted(face):
318                    del_builder.Delete(face)
319    
320    
321class ChamferFilter(FilterSource):
322    name = "Chamfer Filter"
323
324    input = Input
325    
326    size = Float(1.0)
327    
328    _n_edges = Int(4)
329    
330    _low = Int(0)
331    
332    edge_id = Range(low='_low', high='_n_edges')
333    
334    selector = Instance(TNaming.TNaming_Selector)
335    
336    traits_view = View(Item('edge_id'),
337                       'size',
338                       'modified')
339    
340    def _size_changed(self, new_size):
341        self.modified = False
342        self.modified = True
343    
344    @on_trait_change("input, edge_id, label")
345    def on_change_selection(self):
346        new_id = self.edge_id
347        input = self.input
348        label = self.label
349        
350        if not all((input, label)): return
351        
352        input_shape = input.shape
353        
354        sel_label = self.label.FindChild(4)
355        selector = TNaming.TNaming_Selector(sel_label)
356        
357        self.selector = selector
358        
359        topo = Topo(input_shape)
360        self._n_edges = topo.number_of_edges()
361        for i,edge in enumerate(topo.edges()):
362            if i==new_id:
363                selector.Select(edge, input_shape)
364                print "got selection!"
365                break
366        else:
367            print "no selection"
368            
369        self.modified = False
370        self.modified = True
371            
372            
373    def execute(self):
374        input_shape = self.input.shape
375        
376        topo = Topo(input_shape)
377        self._n_edges = topo.number_of_edges()
378        
379        builder = BRepFilletAPI.BRepFilletAPI_MakeChamfer(input_shape)
380        
381        Map = TDF.TDF_LabelMap()
382        itr = TDF.TDF_ChildIterator(self.parent_label, True)
383        while itr.More():
384            sub_label = itr.Value()
385            Map.Add(sub_label)
386            itr.Next()
387            
388        selector = self.selector
389        ret = selector.Solve(Map)
390        
391        if not ret:
392            raise Exception("Failed to solve for edge")
393            #print "Failed to solve for edge"
394        
395        nt = TNaming.TNaming_Tool()
396        selected_shape = nt.CurrentShape(selector.NamedShape())
397        
398        selected_edge = TopoDS.TopoDS().Edge(selected_shape)
399        
400        try:
401            face = Topo(input_shape).faces_from_edge(selected_edge).next()
402        except RuntimeError:
403            raise #not sure how to handle this
404        size = self.size
405        builder.Add(size, size, selected_edge, face)
406        
407        self.update_naming(builder)
408        return builder.Shape()