1#!/usr/bin/env python3
2# Copyright (c) 2020, Matthew Broadway
3# License: MIT License
4import math
5import os
6import time
7from typing import Iterable, Tuple, List, Dict
8
9from PyQt5 import QtWidgets as qw, QtCore as qc, QtGui as qg
10
11import ezdxf
12from ezdxf import recover
13from ezdxf.addons import odafc
14from ezdxf.addons.drawing import Frontend, RenderContext
15
16from ezdxf.addons.drawing.properties import is_dark_color
17from ezdxf.addons.drawing.pyqt import (
18    _get_x_scale, PyQtBackend, CorrespondingDXFEntity,
19    CorrespondingDXFParentStack,
20)
21from ezdxf.audit import Auditor
22from ezdxf.document import Drawing
23from ezdxf.entities import DXFGraphic
24from ezdxf.lldxf.const import DXFStructureError
25
26
27class CADGraphicsView(qw.QGraphicsView):
28    def __init__(self, *, view_buffer: float = 0.2,
29                 zoom_per_scroll_notch: float = 0.2,
30                 loading_overlay: bool = True):
31        super().__init__()
32        self._zoom = 1
33        self._default_zoom = 1
34        self._zoom_limits = (0.5, 100)
35        self._zoom_per_scroll_notch = zoom_per_scroll_notch
36        self._view_buffer = view_buffer
37        self._loading_overlay = loading_overlay
38        self._is_loading = False
39
40        self.setTransformationAnchor(qw.QGraphicsView.AnchorUnderMouse)
41        self.setResizeAnchor(qw.QGraphicsView.AnchorUnderMouse)
42        self.setVerticalScrollBarPolicy(qc.Qt.ScrollBarAlwaysOff)
43        self.setHorizontalScrollBarPolicy(qc.Qt.ScrollBarAlwaysOff)
44        self.setDragMode(qw.QGraphicsView.ScrollHandDrag)
45        self.setFrameShape(qw.QFrame.NoFrame)
46        self.setRenderHints(
47            qg.QPainter.Antialiasing | qg.QPainter.TextAntialiasing | qg.QPainter.SmoothPixmapTransform)
48
49    def clear(self):
50        pass
51
52    def begin_loading(self):
53        self._is_loading = True
54        self.scene().invalidate(qc.QRectF(), qw.QGraphicsScene.AllLayers)
55        qw.QApplication.processEvents()
56
57    def end_loading(self, new_scene: qw.QGraphicsScene):
58        self.setScene(new_scene)
59        self._is_loading = False
60        self.buffer_scene_rect()
61        self.scene().invalidate(qc.QRectF(), qw.QGraphicsScene.AllLayers)
62
63    def buffer_scene_rect(self):
64        scene = self.scene()
65        r = scene.sceneRect()
66        bx, by = r.width() * self._view_buffer / 2, r.height() * self._view_buffer / 2
67        scene.setSceneRect(r.adjusted(-bx, -by, bx, by))
68
69    def fit_to_scene(self):
70        self.fitInView(self.sceneRect(), qc.Qt.KeepAspectRatio)
71        self._default_zoom = _get_x_scale(self.transform())
72        self._zoom = 1
73
74    def _get_zoom_amount(self) -> float:
75        return _get_x_scale(self.transform()) / self._default_zoom
76
77    def wheelEvent(self, event: qg.QWheelEvent) -> None:
78        # dividing by 120 gets number of notches on a typical scroll wheel. See QWheelEvent documentation
79        delta_notches = event.angleDelta().y() / 120
80        direction = math.copysign(1, delta_notches)
81        factor = (1 + self._zoom_per_scroll_notch * direction) ** abs(
82            delta_notches)
83        resulting_zoom = self._zoom * factor
84        if resulting_zoom < self._zoom_limits[0]:
85            factor = self._zoom_limits[0] / self._zoom
86        elif resulting_zoom > self._zoom_limits[1]:
87            factor = self._zoom_limits[1] / self._zoom
88        self.scale(factor, factor)
89        self._zoom *= factor
90
91    def drawForeground(self, painter: qg.QPainter, rect: qc.QRectF) -> None:
92        if self._is_loading and self._loading_overlay:
93            painter.save()
94            painter.fillRect(rect, qg.QColor('#aa000000'))
95            painter.setWorldMatrixEnabled(False)
96            r = self.viewport().rect()
97            painter.setBrush(qc.Qt.NoBrush)
98            painter.setPen(qc.Qt.white)
99            painter.drawText(r.center(), 'Loading...')
100            painter.restore()
101
102
103class CADGraphicsViewWithOverlay(CADGraphicsView):
104    mouse_moved = qc.pyqtSignal(qc.QPointF)
105    element_selected = qc.pyqtSignal(object, int)
106
107    def __init__(self, **kwargs):
108        super().__init__(**kwargs)
109        self._selected_items: List[qw.QGraphicsItem] = []
110        self._selected_index = None
111
112    def clear(self):
113        super().clear()
114        self._selected_items = None
115        self._selected_index = None
116
117    def drawForeground(self, painter: qg.QPainter, rect: qc.QRectF) -> None:
118        super().drawForeground(painter, rect)
119        if self._selected_items:
120            item = self._selected_items[self._selected_index]
121            r = item.sceneTransform().mapRect(item.boundingRect())
122            painter.fillRect(r, qg.QColor(0, 255, 0, 100))
123
124    def mouseMoveEvent(self, event: qg.QMouseEvent) -> None:
125        super().mouseMoveEvent(event)
126        pos = self.mapToScene(event.pos())
127        self.mouse_moved.emit(pos)
128        selected_items = self.scene().items(pos)
129        if selected_items != self._selected_items:
130            self._selected_items = selected_items
131            self._selected_index = 0 if self._selected_items else None
132            self._emit_selected()
133
134    def mouseReleaseEvent(self, event: qg.QMouseEvent) -> None:
135        super().mouseReleaseEvent(event)
136        if event.button() == qc.Qt.LeftButton and self._selected_items:
137            self._selected_index = (self._selected_index + 1) % len(
138                self._selected_items)
139            self._emit_selected()
140
141    def _emit_selected(self):
142        self.element_selected.emit(self._selected_items, self._selected_index)
143        self.scene().invalidate(self.sceneRect(),
144                                qw.QGraphicsScene.ForegroundLayer)
145
146
147class CadViewer(qw.QMainWindow):
148    def __init__(self, params: Dict):
149        super().__init__()
150        self.doc = None
151        self._render_params = params
152        self._render_context = None
153        self._visible_layers = None
154        self._current_layout = None
155        self._reset_backend()
156
157        self.view = CADGraphicsViewWithOverlay()
158        self.view.setScene(qw.QGraphicsScene())
159        self.view.scale(1, -1)  # so that +y is up
160        self.view.element_selected.connect(self._on_element_selected)
161        self.view.mouse_moved.connect(self._on_mouse_moved)
162
163        menu = self.menuBar()
164        select_doc_action = qw.QAction('Select Document', self)
165        select_doc_action.triggered.connect(self._select_doc)
166        menu.addAction(select_doc_action)
167        self.select_layout_menu = menu.addMenu('Select Layout')
168
169        toggle_sidebar_action = qw.QAction('Toggle Sidebar', self)
170        toggle_sidebar_action.triggered.connect(self._toggle_sidebar)
171        menu.addAction(toggle_sidebar_action)
172
173        self.sidebar = qw.QSplitter(qc.Qt.Vertical)
174        self.layers = qw.QListWidget()
175        self.layers.setStyleSheet(
176            'QListWidget {font-size: 12pt;} QCheckBox {font-size: 12pt; padding-left: 5px;}')
177        self.sidebar.addWidget(self.layers)
178        info_container = qw.QWidget()
179        info_layout = qw.QVBoxLayout()
180        info_layout.setContentsMargins(0, 0, 0, 0)
181        self.selected_info = qw.QPlainTextEdit()
182        self.selected_info.setReadOnly(True)
183        info_layout.addWidget(self.selected_info)
184        self.mouse_pos = qw.QLabel()
185        info_layout.addWidget(self.mouse_pos)
186        info_container.setLayout(info_layout)
187        self.sidebar.addWidget(info_container)
188
189        container = qw.QSplitter()
190        self.setCentralWidget(container)
191        container.addWidget(self.view)
192        container.addWidget(self.sidebar)
193        container.setCollapsible(0, False)
194        container.setCollapsible(1, True)
195        w = container.width()
196        container.setSizes([int(3 * w / 4), int(w / 4)])
197
198        self.setWindowTitle('CAD Viewer')
199        self.resize(1600, 900)
200        self.show()
201
202    def _reset_backend(self):
203        # clear caches
204        self._backend = PyQtBackend(use_text_cache=True,
205                                    params=self._render_params)
206
207    def _select_doc(self):
208        path, _ = qw.QFileDialog.getOpenFileName(
209            self, caption='Select CAD Document',
210            filter='CAD Documents (*.dxf *.DXF *.dwg *.DWG)'
211        )
212        if path:
213            try:
214                if os.path.splitext(path)[1].lower() == '.dwg':
215                    doc = odafc.readfile(path)
216                    auditor = doc.audit()
217                else:
218                    try:
219                        doc = ezdxf.readfile(path)
220                    except ezdxf.DXFError:
221                        doc, auditor = recover.readfile(path)
222                    else:
223                        auditor = doc.audit()
224                self.set_document(doc, auditor)
225            except IOError as e:
226                qw.QMessageBox.critical(self, 'Loading Error', str(e))
227            except DXFStructureError as e:
228                qw.QMessageBox.critical(self, 'DXF Structure Error',
229                                        f'Invalid DXF file "{path}": {str(e)}')
230
231    def set_document(self, document: Drawing, auditor: Auditor):
232        error_count = len(auditor.errors)
233        if error_count > 0:
234            ret = qw.QMessageBox.question(
235                self, 'Found DXF Errors',
236                f'Found {error_count} errors in file "{document.filename}"\nLoad file anyway? '
237            )
238            if ret == qw.QMessageBox.No:
239                auditor.print_error_report(auditor.errors)
240                return
241        self.doc = document
242        self._render_context = RenderContext(document)
243        self._reset_backend()  # clear caches
244        self._visible_layers = None
245        self._current_layout = None
246        self._populate_layouts()
247        self._populate_layer_list()
248        self.draw_layout('Model')
249        self.setWindowTitle('CAD Viewer - ' + str(document.filename))
250
251    def _populate_layer_list(self):
252        self.layers.blockSignals(True)
253        self.layers.clear()
254        for layer in self._render_context.layers.values():
255            name = layer.layer
256            item = qw.QListWidgetItem()
257            self.layers.addItem(item)
258            checkbox = qw.QCheckBox(name)
259            checkbox.setCheckState(
260                qc.Qt.Checked if layer.is_visible else qc.Qt.Unchecked)
261            checkbox.stateChanged.connect(self._layers_updated)
262            text_color = '#FFFFFF' if is_dark_color(layer.color,
263                                                    0.4) else '#000000'
264            checkbox.setStyleSheet(
265                f'color: {text_color}; background-color: {layer.color}')
266            self.layers.setItemWidget(item, checkbox)
267        self.layers.blockSignals(False)
268
269    def _populate_layouts(self):
270        self.select_layout_menu.clear()
271        for layout_name in self.doc.layout_names_in_taborder():
272            action = qw.QAction(layout_name, self)
273            action.triggered.connect(
274                lambda: self.draw_layout(layout_name, reset_view=True))
275            self.select_layout_menu.addAction(action)
276
277    def draw_layout(self, layout_name: str, reset_view: bool = True):
278        print(f'drawing {layout_name}')
279        self._current_layout = layout_name
280        self.view.begin_loading()
281        new_scene = qw.QGraphicsScene()
282        self._backend.set_scene(new_scene)
283        layout = self.doc.layout(layout_name)
284        self._update_render_context(layout)
285        try:
286            start = time.perf_counter()
287            Frontend(self._render_context, self._backend).draw_layout(layout)
288            duration = time.perf_counter() - start
289            print(f'took {duration:.4f} seconds')
290        except DXFStructureError as e:
291            qw.QMessageBox.critical(
292                self, 'DXF Structure Error',
293                f'Abort rendering of layout "{layout_name}": {str(e)}')
294        finally:
295            self._backend.finalize()
296        self.view.end_loading(new_scene)
297        self.view.buffer_scene_rect()
298        if reset_view:
299            self.view.fit_to_scene()
300
301    def _update_render_context(self, layout):
302        assert self._render_context
303        self._render_context.set_current_layout(layout)
304        # Direct modification of RenderContext.layers would be more flexible, but would also expose the internals.
305        if self._visible_layers is not None:
306            self._render_context.set_layers_state(self._visible_layers,
307                                                  state=True)
308
309    def resizeEvent(self, event: qg.QResizeEvent) -> None:
310        self.view.fit_to_scene()
311
312    def _layer_checkboxes(self) -> Iterable[Tuple[int, qw.QCheckBox]]:
313        for i in range(self.layers.count()):
314            item = self.layers.itemWidget(self.layers.item(i))
315            yield i, item
316
317    @qc.pyqtSlot(int)
318    def _layers_updated(self, item_state: qc.Qt.CheckState):
319        shift_held = qw.QApplication.keyboardModifiers() & qc.Qt.ShiftModifier
320        if shift_held:
321            for i, item in self._layer_checkboxes():
322                item.blockSignals(True)
323                item.setCheckState(item_state)
324                item.blockSignals(False)
325
326        self._visible_layers = set()
327        for i, layer in self._layer_checkboxes():
328            if layer.checkState() == qc.Qt.Checked:
329                self._visible_layers.add(layer.text())
330        self.draw_layout(self._current_layout, reset_view=False)
331
332    @qc.pyqtSlot()
333    def _toggle_sidebar(self):
334        self.sidebar.setHidden(not self.sidebar.isHidden())
335
336    @qc.pyqtSlot(qc.QPointF)
337    def _on_mouse_moved(self, mouse_pos: qc.QPointF):
338        self.mouse_pos.setText(
339            f'mouse position: {mouse_pos.x():.4f}, {mouse_pos.y():.4f}\n')
340
341    @qc.pyqtSlot(object, int)
342    def _on_element_selected(self, elements: List[qw.QGraphicsItem],
343                             index: int):
344        if not elements:
345            text = 'No element selected'
346        else:
347            text = f'Selected: {index + 1} / {len(elements)}    (click to cycle)\n'
348            element = elements[index]
349            dxf_entity = element.data(CorrespondingDXFEntity)
350            if dxf_entity is None:
351                text += 'No data'
352            else:
353                text += f'Selected Entity: {dxf_entity}\nLayer: {dxf_entity.dxf.layer}\n\nDXF Attributes:\n'
354                text += _entity_attribs_string(dxf_entity)
355
356                dxf_parent_stack = element.data(CorrespondingDXFParentStack)
357                if dxf_parent_stack:
358                    text += '\nParents:\n'
359                    for entity in reversed(dxf_parent_stack):
360                        text += f'- {entity}\n'
361                        text += _entity_attribs_string(entity, indent='    ')
362
363        self.selected_info.setPlainText(text)
364
365
366def _entity_attribs_string(dxf_entity: DXFGraphic, indent: str = '') -> str:
367    text = ''
368    for key, value in dxf_entity.dxf.all_existing_dxf_attribs().items():
369        text += f'{indent}- {key}: {value}\n'
370    return text
371