Add 'point_to_row' function to avoid code duplication

scripts/dxf_to_csv.py

@@ -17,9 +17,20 @@ import sys
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
+import argparse
 import dxfgrabber
 
 
+def point_to_row(point, layer, dxf_type, dxf_id):
+    """Create a dict from a point with layer and e properties"""
+    row = {key: val for key, val in zip(['x', 'y', 'z'], [*point])}
+    row['layer'] = layer
+    row['type'] = dxf_type
+    row['id'] = dxf_id
+
+    return row
+
+
 def convert(dxf_name, layer_name):
     # Read dxf file
     dxf = dxfgrabber.readfile(dxf_name)
@@ -42,39 +53,26 @@ def convert(dxf_name, layer_name):
 
     # Get coordinates of entities
     data = []
-    for entity in entities:
-        if entity.dxftype == 'POINT':
-            point = entity.point
-            row = {key: val for key, val in zip(['x', 'y', 'z'], [*point])}
-            row['layer'] = entity.layer
-            row['id'] = entity.handle
-            row['type'] = entity.dxftype
+    for e in entities:
+        if e.dxftype == 'POINT':
+            point = e.point
+            row = point_to_row(point, e.layer, e.dxftype, e.handle)
             data.append(row)
-        elif entity.dxftype == 'LINE':
-            points = [entity.start, entity.end]
+        elif e.dxftype == 'LINE':
+            points = [e.start, e.end]
             for point in points:
-                row = {key: val for key, val in zip(['x', 'y', 'z'], [*point])}
-                row['layer'] = entity.layer
-                row['id'] = entity.handle
-                row['type'] = entity.dxftype
+                row = point_to_row(point, e.layer, e.dxftype, e.handle)
                 data.append(row)
-        elif entity.dxftype == 'POLYLINE':
-            for vertex in entity.vertices:
+        elif e.dxftype == 'POLYLINE':
+            for vertex in e.vertices:
                 point = vertex.location
-                row = {key: val for key, val in zip(['x', 'y', 'z'], [*point])}
-                row['layer'] = entity.layer
-                row['id'] = entity.handle
-                row['type'] = entity.dxftype
+                row = point_to_row(point, e.layer, e.dxftype, e.handle)
                 data.append(row)
-        elif entity.dxftype == 'LWPOLYLINE':
-            for point in entity.points:
-                row = {key: val for key, val in zip(['x', 'y', 'z'], [*point])}
-                row['layer'] = entity.layer
-                row['id'] = entity.handle
-                row['type'] = entity.dxftype
+        elif e.dxftype == 'LWPOLYLINE':
+            for point in e.points:
+                row = point_to_row(point, e.layer, e.dxftype, e.handle)
                 data.append(row)
 
-
     # Add to dataframe
     points = pd.DataFrame(data)
     points = points[['layer', 'type', 'id', 'x', 'y', 'z']]