Commits

Chris Mutel  committed 051f59b

Allow imports of directories or files

  • Participants
  • Parent commits dd974cf

Comments (0)

Files changed (2)

File brightway2/io/import_ecospold.py

 from lxml import objectify
 import math
 import os
+import progressbar
 from stats_toolkit.distributions import *
-try:
-    import progressbar
-except ImportError:
-    progressbar = None
 
 BIOSPHERE = ("air", "water", "soil", "resource")
 
 
 class EcospoldImporter(object):
-    def import_directory(self, path, name, depends=["biosphere", ]):
+    def importer(self, path, name, depends=["biosphere", ]):
         data = []
         log = get_logger(name)
         log.critical(u"Starting import of %s (from %s)" % (name, path))
-        files = filter(lambda x: x[-4:].lower() == ".xml", os.listdir(path))
+        if os.path.isdir(path):
+            files = [os.path.join(path, filter(lambda x: x[-4:].lower(
+                ) == ".xml", os.listdir(path)))]
+        else:
+            files = [path]
 
-        if progressbar:
-            widgets = ['Files: ', progressbar.Percentage(), ' ',
-                progressbar.Bar(marker=progressbar.RotatingMarker()), ' ',
-                progressbar.ETA()]
-            pbar = progressbar.ProgressBar(widgets=widgets, maxval=len(files)
-                ).start()
+        widgets = ['Files: ', progressbar.Percentage(), ' ',
+            progressbar.Bar(marker=progressbar.RotatingMarker()), ' ',
+            progressbar.ETA()]
+        pbar = progressbar.ProgressBar(widgets=widgets, maxval=len(files)
+            ).start()
 
         for index, filename in enumerate(files):
-            root = objectify.parse(open(os.path.join(path, filename))
-                ).getroot()
+            root = objectify.parse(open(filename)).getroot()
 
             if root.tag != '{http://www.EcoInvent.org/EcoSpold01}ecoSpold':
                 # Unrecognized file type
             for dataset in root.iterchildren():
                 data.append(self._process_dataset(dataset))
 
-            if progressbar:
-                pbar.update(index)
-
-        if progressbar:
-            pbar.finish()
+            pbar.update(index)
 
         # Hackish
         for o in data:
         manager.register("Ecospold 1", depends, len(data))
         manager.write(data)
 
+        pbar.finish()
+
     def _find_in_dependent_database(self, code, exc, depends):
         for db in depends:
             if (db, code) in mapping:
                 continue
 
             this = {
-                "pedigree matrix": exc.get("generalComment"),
                 "code": int(exc.get("number")),
                 "_matching": {
                     "categories": (exc.get("category"), exc.get("subCategory")),
                     }
                 }
 
+            if exc.get("generalComment"):
+                this["pedigree matrix"] = exc.get("generalComment")
+
             uncertainty = int(exc.get("uncertaintyType", 0))
             mean = exc.get("meanValue")
             min_ = exc.get("minValue")

File brightway2/io/import_method.py

 from brightway2 import Database, mapping, Method, methods
 from lxml import objectify
 import os
+import progressbar
 try:
     import cPickle as pickle
 except:
     import pickle
 
 
-def import_ia_dir(dirpath):
-    for filename in filter(lambda x: x.lower()[-4:] == ".xml",
-            os.listdir(dirpath)):
-        filepath = os.path.join(dirpath, filename)
-        print "Working on %s" % filepath
-        EcospoldImpactAssessmentImporter(filepath)
-
-
 class EcospoldImpactAssessmentImporter(object):
     """
 Import impact assessment methods and weightings from ecospold XML format.
     """
-    def __init__(self, filename):
-        self.filename = filename
+    def importer(self, path):
+        if os.path.isdir(path):
+            files = [os.path.join(path, filter(lambda x: x[-4:].lower(
+                ) == ".xml", os.listdir(path)))]
+        else:
+            files = [path]
+
         self.biosphere_data = Database("biosphere").load()
-        # Note that this is only used for the first root method found in
-        # the file
-        root = objectify.parse(open(self.filename)).getroot()
-        for dataset in root.iterchildren():
-            self.add_method(dataset)
+        if progressbar:
+            widgets = ['Files: ', progressbar.Percentage(), ' ',
+                progressbar.Bar(marker=progressbar.RotatingMarker()), ' ',
+                progressbar.ETA()]
+            pbar = progressbar.ProgressBar(widgets=widgets, maxval=len(files)
+                ).start()
+
+        for index, filepath in enumerate(files):
+            # Note that this is only used for the first root method found in
+            # the file
+            root = objectify.parse(open(filepath)).getroot()
+            for dataset in root.iterchildren():
+                self.add_method(dataset)
+            pbar.update(index)
+
+        pbar.finish()
 
     def add_method(self, ds):
         ref_func = ds.metaInformation.processInformation.referenceFunction