From 2896c231434cfdb0ad75996249e370f803f169ee Mon Sep 17 00:00:00 2001 From: Dion Moult Date: Wed, 7 Jun 2023 15:59:20 +1000 Subject: [PATCH] Add support for Pandas DataFrames in IfcCSV and clean up function signatures --- .../blenderbim/bim/module/csv/operator.py | 19 ++--- src/ifccsv/ifccsv.py | 83 ++++++++++--------- 2 files changed, 49 insertions(+), 53 deletions(-) diff --git a/src/blenderbim/blenderbim/bim/module/csv/operator.py b/src/blenderbim/blenderbim/bim/module/csv/operator.py index 07a2206795..02b7fabe41 100644 --- a/src/blenderbim/blenderbim/bim/module/csv/operator.py +++ b/src/blenderbim/blenderbim/bim/module/csv/operator.py @@ -150,14 +150,9 @@ class ExportIfcCsv(bpy.types.Operator): selector = ifcopenshell.util.selector.Selector() results = selector.parse(ifc_file, props.ifc_selector) ifc_csv = ifccsv.IfcCsv() - ifc_csv.output = self.filepath - ifc_csv.attributes = [a.name for a in props.csv_attributes] - ifc_csv.selector = selector - if props.csv_delimiter == "CUSTOM": - ifc_csv.delimiter = props.csv_custom_delimiter - else: - ifc_csv.delimiter = props.csv_delimiter - ifc_csv.export(ifc_file, results) + attributes = [a.name for a in props.csv_attributes] + sep = props.csv_custom_delimiter if props.csv_delimiter == "CUSTOM" else props.csv_delimiter + ifc_csv.export(ifc_file, results, attributes, output=self.filepath, format=args.format, delimiter=sep) return {"FINISHED"} @@ -182,12 +177,8 @@ class ImportIfcCsv(bpy.types.Operator): else: ifc_file = ifcopenshell.open(props.csv_ifc_file) ifc_csv = ifccsv.IfcCsv() - ifc_csv.output = self.filepath - if props.csv_delimiter == "CUSTOM": - ifc_csv.delimiter = props.csv_custom_delimiter - else: - ifc_csv.delimiter = props.csv_delimiter - ifc_csv.Import(ifc_file) + sep = props.csv_custom_delimiter if props.csv_delimiter == "CUSTOM" else props.csv_delimiter + ifc_csv.Import(ifc_file, self.filepath, delimiter=sep) if not props.should_load_from_memory: ifc_file.write(props.csv_ifc_file) purge_module_data() diff --git a/src/ifccsv/ifccsv.py b/src/ifccsv/ifccsv.py index e415f4bad2..6497fe3671 100755 --- a/src/ifccsv/ifccsv.py +++ b/src/ifccsv/ifccsv.py @@ -42,6 +42,12 @@ except: pass # No XLSX support +try: + import pandas as pd +except: + pass # No Pandas support + + class IfcAttributeSetter: @staticmethod def set_element_key(ifc_file, element, key, value): @@ -139,15 +145,14 @@ class IfcAttributeSetter: class IfcCsv: - def __init__(self): + def __init__(self, output="", delimiter=","): + self.headers = [] self.results = [] - self.attributes = [] - self.output = "" - self.format = "csv" - self.delimiter = "," + self.dataframe = None - def export(self, ifc_file, elements): + def export(self, ifc_file, elements, attributes, output=None, format=None, delimiter=","): self.ifc_file = ifc_file + self.results = [] for element in elements: result = [] if hasattr(element, "GlobalId"): @@ -155,33 +160,35 @@ class IfcCsv: else: result.append(None) - for index, attribute in enumerate(self.attributes): + for index, attribute in enumerate(attributes or []): if "*" in attribute: - self.attributes.extend(self.get_wildcard_attributes(attribute)) - del self.attributes[index] + attributes.extend(self.get_wildcard_attributes(attribute)) + del attributes[index] - for attribute in self.attributes: + for attribute in attributes: result.append(ifcopenshell.util.selector.get_element_value(element, attribute)) self.results.append(result) self.headers = ["GlobalId"] - self.headers.extend(self.attributes) + self.headers.extend(attributes or []) - if self.format == "csv": - self.export_csv() - elif self.format == "ods": - self.export_ods() - elif self.format == "xlsx": - self.export_xlsx() + if format == "csv": + self.export_csv(output, delimiter=delimiter) + elif format == "ods": + self.export_ods(output) + elif format == "xlsx": + self.export_xlsx(output) + elif format == "pd": + return self.export_pd() - def export_csv(self): - with open(self.output, "w", newline="", encoding="utf-8") as f: - writer = csv.writer(f, delimiter=self.delimiter) + def export_csv(self, output, delimiter=None): + with open(output, "w", newline="", encoding="utf-8") as f: + writer = csv.writer(f, delimiter=delimiter) writer.writerow(self.headers) for row in self.results: writer.writerow(row) - def export_ods(self): + def export_ods(self, output): self.doc = OpenDocumentSpreadsheet() self.colours = { @@ -218,12 +225,12 @@ class IfcCsv: table.addElement(tr) self.doc.spreadsheet.addElement(table) - if self.output[-4:].lower() == ".ods": - self.output = self.output[0:-4] - self.doc.save(self.output, True) + if output[-4:].lower() == ".ods": + output = output[0:-4] + self.doc.save(output, True) - def export_xlsx(self): - self.workbook = Workbook(self.output) + def export_xlsx(self, output): + self.workbook = Workbook(output) self.colours = { "h": "dc8774", # Header @@ -254,6 +261,10 @@ class IfcCsv: self.workbook.close() + def export_pd(self): + self.dataframe = pd.DataFrame(self.results, columns=self.headers) + return self.dataframe + def get_wildcard_attributes(self, attribute): results = set() pset_qto_name = attribute.split(".", 1)[0] @@ -266,9 +277,10 @@ class IfcCsv: results.update([p.Name for p in element.Quantities]) return ["{}.{}".format(pset_qto_name, n) for n in results] - def Import(self, ifc_file): - with open(self.output, newline="", encoding="utf-8") as f: - reader = csv.reader(f, delimiter=self.delimiter) + def Import(self, ifc_file, table, delimiter=","): + # Currently only supports CSV. + with open(table, newline="", encoding="utf-8") as f: + reader = csv.reader(f, delimiter=delimiter) headers = [] for row in reader: if not headers: @@ -303,18 +315,11 @@ if __name__ == "__main__": if args.export: ifc_file = ifcopenshell.open(args.ifc) - selector = ifcopenshell.util.selector.Selector() - results = selector.parse(ifc_file, args.query) + results = ifcopenshell.util.selector.Selector.parse(ifc_file, args.query) ifc_csv = IfcCsv() - ifc_csv.output = args.spreadsheet - ifc_csv.format = args.format - ifc_csv.attributes = args.arguments if args.arguments else [] - ifc_csv.selector = selector - ifc_csv.export(ifc_file, results) + ifc_csv.export(ifc_file, results, args.arguments or [], output=args.spreadsheet, format=args.format) elif getattr(args, "import"): ifc_csv = IfcCsv() - ifc_csv.output = args.spreadsheet - ifc_csv.format = args.format ifc_file = ifcopenshell.open(args.ifc) - ifc_csv.Import(ifc_file) + ifc_csv.Import(ifc_file, args.spreadsheet) ifc_file.write(args.ifc)