1# created: 2019-01-03
2# Copyright (c) 2019-2020 Manfred Moitzi
3# License: MIT License
4from typing import TYPE_CHECKING, Iterable, Dict
5from ezdxf.math import Vec2, Shape2d, NULLVEC
6from .forms import open_arrow, arrow2
7
8if TYPE_CHECKING:
9    from ezdxf.eztypes import Vertex, GenericLayoutType, DXFGraphic, Drawing
10
11DEFAULT_ARROW_ANGLE = 18.924644
12DEFAULT_BETA = 45.
13
14
15# The base arrow is oriented for the right hand side ->| of the dimension line, reverse is the left hand side |<-.
16class BaseArrow:
17    def __init__(self, vertices: Iterable['Vertex']):
18        self.shape = Shape2d(vertices)
19
20    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
21        pass
22
23    def place(self, insert: 'Vertex', angle: float):
24        self.shape.rotate(angle)
25        self.shape.translate(insert)
26
27
28class NoneStroke(BaseArrow):
29    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
30        super().__init__([Vec2(insert)])
31
32
33class ObliqueStroke(BaseArrow):
34    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
35        self.size = size
36        s2 = size / 2
37        # shape = [center, lower left, upper right]
38        super().__init__([Vec2((-s2, -s2)), Vec2((s2, s2))])
39        self.place(insert, angle)
40
41    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
42        layout.add_line(start=self.shape[0], end=self.shape[1], dxfattribs=dxfattribs)
43
44
45class ArchTick(ObliqueStroke):
46    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
47        width = self.size * .15
48        if layout.dxfversion > 'AC1009':
49            dxfattribs['const_width'] = width
50            layout.add_lwpolyline(self.shape, format='xy', dxfattribs=dxfattribs)
51        else:
52            dxfattribs['default_start_width'] = width
53            dxfattribs['default_end_width'] = width
54            layout.add_polyline2d(self.shape, dxfattribs=dxfattribs)
55
56
57class ClosedArrowBlank(BaseArrow):
58    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
59        super().__init__(open_arrow(size, angle=DEFAULT_ARROW_ANGLE))
60        self.place(insert, angle)
61
62    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
63        if layout.dxfversion > 'AC1009':
64            polyline = layout.add_lwpolyline(
65                points=self.shape,
66                dxfattribs=dxfattribs)
67        else:
68            polyline = layout.add_polyline2d(
69                points=self.shape,
70                dxfattribs=dxfattribs)
71        polyline.close(True)
72
73
74class ClosedArrow(ClosedArrowBlank):
75    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
76        super().render(layout, dxfattribs)
77        end_point = self.shape[0].lerp(self.shape[2])
78
79        layout.add_line(start=self.shape[1], end=end_point, dxfattribs=dxfattribs)
80
81
82class ClosedArrowFilled(ClosedArrow):
83    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
84        layout.add_solid(
85            points=self.shape,
86            dxfattribs=dxfattribs,
87        )
88
89
90class _OpenArrow(BaseArrow):
91    def __init__(self, arrow_angle: float, insert: 'Vertex', size: float = 1.0, angle: float = 0):
92        points = list(open_arrow(size, angle=arrow_angle))
93        points.append((-1, 0))
94        super().__init__(points)
95        self.place(insert, angle)
96
97    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
98        if layout.dxfversion > 'AC1009':
99            layout.add_lwpolyline(points=self.shape[:-1], dxfattribs=dxfattribs)
100        else:
101            layout.add_polyline2d(points=self.shape[:-1], dxfattribs=dxfattribs)
102        layout.add_line(start=self.shape[1], end=self.shape[-1], dxfattribs=dxfattribs)
103
104
105class OpenArrow(_OpenArrow):
106    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
107        super().__init__(DEFAULT_ARROW_ANGLE, insert, size, angle)
108
109
110class OpenArrow30(_OpenArrow):
111    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
112        super().__init__(30, insert, size, angle)
113
114
115class OpenArrow90(_OpenArrow):
116    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
117        super().__init__(90, insert, size, angle)
118
119
120class Circle(BaseArrow):
121    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
122        self.radius = size / 2
123        # shape = [center point, connection point]
124        super().__init__([
125            Vec2((0, 0)),
126            Vec2((-self.radius, 0)),
127            Vec2((-size, 0)),
128        ])
129        self.place(insert, angle)
130
131    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
132        layout.add_circle(center=self.shape[0], radius=self.radius, dxfattribs=dxfattribs)
133
134
135class Origin(Circle):
136    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
137        super().render(layout, dxfattribs)
138        layout.add_line(start=self.shape[0], end=self.shape[2], dxfattribs=dxfattribs)
139
140
141class CircleBlank(Circle):
142    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
143        super().render(layout, dxfattribs)
144        layout.add_line(start=self.shape[1], end=self.shape[2], dxfattribs=dxfattribs)
145
146
147class Origin2(Circle):
148    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
149        layout.add_circle(center=self.shape[0], radius=self.radius, dxfattribs=dxfattribs)
150        layout.add_circle(center=self.shape[0], radius=self.radius / 2, dxfattribs=dxfattribs)
151        layout.add_line(start=self.shape[1], end=self.shape[2], dxfattribs=dxfattribs)
152
153
154class DotSmall(Circle):
155    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
156        center = self.shape[0]
157        d = Vec2((self.radius / 2, 0))
158        p1 = center - d
159        p2 = center + d
160        if layout.dxfversion > 'AC1009':
161            dxfattribs['const_width'] = self.radius
162            layout.add_lwpolyline([(p1, 1), (p2, 1)], format='vb', close=True,
163                                  dxfattribs=dxfattribs)
164        else:
165            dxfattribs['default_start_width'] = self.radius
166            dxfattribs['default_end_width'] = self.radius
167            polyline = layout.add_polyline2d(points=[p1, p2], close=True,
168                                             dxfattribs=dxfattribs)
169            polyline[0].dxf.bulge = 1
170            polyline[1].dxf.bulge = 1
171
172
173class Dot(DotSmall):
174    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
175        layout.add_line(start=self.shape[1], end=self.shape[2], dxfattribs=dxfattribs)
176        super().render(layout, dxfattribs)
177
178
179class Box(BaseArrow):
180    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
181        # shape = [lower_left, lower_right, upper_right, upper_left, connection point]
182        s2 = size / 2
183        super().__init__([
184            Vec2((-s2, -s2)),
185            Vec2((+s2, -s2)),
186            Vec2((+s2, +s2)),
187            Vec2((-s2, +s2)),
188            Vec2((-s2, 0)),
189            Vec2((-size, 0)),
190        ])
191        self.place(insert, angle)
192
193    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
194        if layout.dxfversion > 'AC1009':
195            polyline = layout.add_lwpolyline(points=self.shape[0:4], dxfattribs=dxfattribs)
196        else:
197            polyline = layout.add_polyline2d(points=self.shape[0:4], dxfattribs=dxfattribs)
198        polyline.close(True)
199        layout.add_line(start=self.shape[4], end=self.shape[5], dxfattribs=dxfattribs)
200
201
202class BoxFilled(Box):
203    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
204        def solid_order():
205            v = self.shape.vertices
206            return [v[0], v[1], v[3], v[2]]
207
208        layout.add_solid(points=solid_order(), dxfattribs=dxfattribs)
209        layout.add_line(start=self.shape[4], end=self.shape[5], dxfattribs=dxfattribs)
210
211
212class Integral(BaseArrow):
213    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
214        self.radius = size * .3535534
215        self.angle = angle
216        # shape = [center, left_center, right_center]
217        super().__init__([
218            Vec2((0, 0)),
219            Vec2((-self.radius, 0)),
220            Vec2((self.radius, 0)),
221        ])
222        self.place(insert, angle)
223
224    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
225        angle = self.angle
226        layout.add_arc(center=self.shape[1], radius=self.radius, start_angle=-90 + angle, end_angle=angle,
227                       dxfattribs=dxfattribs)
228        layout.add_arc(center=self.shape[2], radius=self.radius, start_angle=90 + angle, end_angle=180 + angle,
229                       dxfattribs=dxfattribs)
230
231
232class DatumTriangle(BaseArrow):
233    REVERSE_ANGLE = 180
234
235    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
236        d = .577350269 * size  # tan(30)
237        # shape = [upper_corner, lower_corner, connection_point]
238        super().__init__([
239            Vec2((0, d)),
240            Vec2((0, -d)),
241            Vec2((-size, 0)),
242        ])
243        self.place(insert, angle)
244
245    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
246        if layout.dxfversion > 'AC1009':
247            polyline = layout.add_lwpolyline(points=self.shape, dxfattribs=dxfattribs)
248        else:
249            polyline = layout.add_polyline2d(points=self.shape, dxfattribs=dxfattribs)
250        polyline.close(True)
251
252
253class DatumTriangleFilled(DatumTriangle):
254    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
255        layout.add_solid(points=self.shape, dxfattribs=dxfattribs)
256
257
258class _EzArrow(BaseArrow):
259    def __init__(self, insert: 'Vertex', size: float = 1.0, angle: float = 0):
260        points = list(arrow2(size, angle=DEFAULT_ARROW_ANGLE))
261        points.append((-1, 0))
262        super().__init__(points)
263        self.place(insert, angle)
264
265    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
266        if layout.dxfversion > 'AC1009':
267            polyline = layout.add_lwpolyline(self.shape[:-1], dxfattribs=dxfattribs)
268        else:
269            polyline = layout.add_polyline2d(self.shape[:-1], dxfattribs=dxfattribs)
270        polyline.close(True)
271
272
273class EzArrowBlank(_EzArrow):
274    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
275        super().render(layout, dxfattribs)
276        layout.add_line(start=self.shape[-2], end=self.shape[-1], dxfattribs=dxfattribs)
277
278
279class EzArrow(_EzArrow):
280    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
281        super().render(layout, dxfattribs)
282        layout.add_line(start=self.shape[1], end=self.shape[-1], dxfattribs=dxfattribs)
283
284
285class EzArrowFilled(_EzArrow):
286    def render(self, layout: 'GenericLayoutType', dxfattribs: dict = None):
287        points = self.shape.vertices
288        layout.add_solid([points[0], points[1], points[3], points[2]], dxfattribs=dxfattribs)
289        layout.add_line(start=self.shape[-2], end=self.shape[-1], dxfattribs=dxfattribs)
290
291
292class _Arrows:
293    closed_filled = ""
294    dot = "DOT"
295    dot_small = "DOTSMALL"
296    dot_blank = "DOTBLANK"
297    origin_indicator = "ORIGIN"
298    origin_indicator_2 = "ORIGIN2"
299    open = "OPEN"
300    right_angle = "OPEN90"
301    open_30 = "OPEN30"
302    closed = "CLOSED"
303    dot_smallblank = "SMALL"
304    none = "NONE"
305    oblique = "OBLIQUE"
306    box_filled = "BOXFILLED"
307    box = "BOXBLANK"
308    closed_blank = "CLOSEDBLANK"
309    datum_triangle_filled = "DATUMFILLED"
310    datum_triangle = "DATUMBLANK"
311    integral = "INTEGRAL"
312    architectural_tick = "ARCHTICK"
313    # ezdxf special arrows
314    ez_arrow = "EZ_ARROW"
315    ez_arrow_blank = "EZ_ARROW_BLANK"
316    ez_arrow_filled = "EZ_ARROW_FILLED"
317
318    CLASSES = {
319        closed_filled: ClosedArrowFilled,
320        dot: Dot,
321        dot_small: DotSmall,
322        dot_blank: CircleBlank,
323        origin_indicator: Origin,
324        origin_indicator_2: Origin2,
325        open: OpenArrow,
326        right_angle: OpenArrow90,
327        open_30: OpenArrow30,
328        closed: ClosedArrow,
329        dot_smallblank: Circle,
330        none: NoneStroke,
331        oblique: ObliqueStroke,
332        box_filled: BoxFilled,
333        box: Box,
334        closed_blank: ClosedArrowBlank,
335        datum_triangle: DatumTriangle,
336        datum_triangle_filled: DatumTriangleFilled,
337        integral: Integral,
338        architectural_tick: ArchTick,
339        ez_arrow: EzArrow,
340        ez_arrow_blank: EzArrowBlank,
341        ez_arrow_filled: EzArrowFilled,
342    }
343    # arrows with origin at dimension line start/end
344    ORIGIN_ZERO = {
345        architectural_tick,
346        oblique,
347        dot_small,
348        dot_smallblank,
349        integral,
350        none,
351    }
352
353    __acad__ = {
354        closed_filled, dot, dot_small, dot_blank, origin_indicator, origin_indicator_2, open, right_angle, open_30,
355        closed, dot_smallblank, none, oblique, box_filled, box, closed_blank, datum_triangle, datum_triangle_filled,
356        integral, architectural_tick
357    }
358    __ezdxf__ = {
359        ez_arrow,
360        ez_arrow_blank,
361        ez_arrow_filled,
362    }
363    __all_arrows__ = __acad__ | __ezdxf__
364
365    EXTENSIONS_ALLOWED = {
366        architectural_tick,
367        oblique,
368        none,
369        dot_smallblank,
370        integral,
371        dot_small,
372    }
373
374    def is_acad_arrow(self, item: str) -> bool:
375        return item.upper() in self.__acad__
376
377    def is_ezdxf_arrow(self, item: str) -> bool:
378        return item.upper() in self.__ezdxf__
379
380    def has_extension_line(self, name):
381        return name in self.EXTENSIONS_ALLOWED
382
383    def __contains__(self, item: str) -> bool:
384        if item is None:
385            return False
386        return item.upper() in self.__all_arrows__
387
388    def create_block(self, blocks, name: str):
389        block_name = self.block_name(name)
390        if block_name not in blocks:
391            block = blocks.new(block_name)
392            arrow = self.arrow_shape(name, insert=(0, 0), size=1, rotation=0)
393            arrow.render(block, dxfattribs={'color': 0, 'linetype': 'BYBLOCK'})
394        return block_name
395
396    def block_name(self, name):
397        if not self.is_acad_arrow(name):  # common BLOCK definition
398            return name.upper()  # e.g. Dimension.dxf.bkl = 'EZ_ARROW' == Insert.dxf.name
399        elif name == '':  # special AutoCAD arrow symbol 'CLOSED_FILLED' has no name
400            # ezdxf uses blocks for ALL arrows, but '_' (closed filled) as block name?
401            return '_CLOSEDFILLED'  # Dimension.dxf.bkl = '' != Insert.dxf.name = '_CLOSED_FILLED'
402        else:  # add preceding '_' to AutoCAD arrow symbol names
403            return '_' + name.upper()  # Dimension.dxf.bkl = 'DOT' != Insert.dxf.name = '_DOT'
404
405    def arrow_name(self, block_name: str) -> str:
406        if block_name.startswith('_'):
407            name = block_name[1:].upper()
408            if name == 'CLOSEDFILLED':
409                return ''
410            elif self.is_acad_arrow(name):
411                return name
412        return block_name
413
414    def insert_arrow(self, layout: 'GenericLayoutType',
415                     name: str,
416                     insert: 'Vertex' = NULLVEC,
417                     size: float = 1.0,
418                     rotation: float = 0, *,
419                     dxfattribs: Dict = None) -> Vec2:
420        """ Insert arrow as block reference into `layout`. """
421        block_name = self.create_block(layout.doc.blocks, name)
422
423        dxfattribs = dict(dxfattribs) if dxfattribs else {}  # copy attribs
424        dxfattribs['rotation'] = rotation
425        dxfattribs['xscale'] = size
426        dxfattribs['yscale'] = size
427        layout.add_blockref(block_name, insert=insert, dxfattribs=dxfattribs)
428        return connection_point(name, insert=insert, scale=size, rotation=rotation)
429
430    def render_arrow(self, layout: 'GenericLayoutType',
431                     name: str,
432                     insert: 'Vertex' = NULLVEC,
433                     size: float = 1.0,
434                     rotation: float = 0, *,
435                     dxfattribs: Dict = None) -> Vec2:
436        """ Render arrow as basic DXF entities into `layout`. """
437        dxfattribs = dxfattribs or {}
438        arrow = self.arrow_shape(name, insert, size, rotation)
439        arrow.render(layout, dxfattribs)
440        return connection_point(name, insert=insert, scale=size, rotation=rotation)
441
442    def virtual_entities(self,
443                         name: str,
444                         insert: 'Vertex' = NULLVEC,
445                         size: float = 0.625,
446                         rotation: float = 0, *,
447                         dxfattribs: Dict = None) -> Iterable['DXFGraphic']:
448        """ Yield arrow components as virtual DXF entities. """
449        from ezdxf.layouts import VirtualLayout
450        if name in self:
451            layout = VirtualLayout()
452            dxfattribs = dxfattribs or {}
453            ARROWS.render_arrow(
454                layout, name,
455                insert=insert,
456                size=size,
457                rotation=rotation,
458                dxfattribs=dxfattribs,
459            )
460            yield from iter(layout)
461
462    def arrow_shape(self, name: str, insert: 'Vertex', size: float, rotation: float) -> BaseArrow:
463        # size depending shapes
464        name = name.upper()
465        if name == self.dot_small:
466            size *= .25
467        elif name == self.dot_smallblank:
468            size *= .5
469        cls = self.CLASSES[name]
470        return cls(insert, size, rotation)
471
472
473def connection_point(arrow_name: str, insert: 'Vertex', scale: float = 1, rotation: float = 0) -> Vec2:
474    insert = Vec2(insert)
475    if arrow_name in _Arrows.ORIGIN_ZERO:
476        return insert
477    else:
478        return insert - Vec2.from_deg_angle(rotation, scale)
479
480
481ARROWS = _Arrows()
482