o Going more jupyter.

1 files changed, 72 insertions, 5 deletions

diff --git a/src/ee/kicad/bom/__init__.py b/src/ee/kicad/bom/__init__.py
index 1808357..d0798a4 100644
--- a/src/ee/kicad/bom/__init__.py
+++ b/src/ee/kicad/bom/__init__.py
@@ -1,10 +1,34 @@
+import re
+import sys
+
 __all__ = [
   'Part',
   'Library',
   'Bom',
   'Comp',
+  'split_ref',
 ]
 
+def split_ref(ref):
+  """Split "C12" into a tuple that's useful for sorting by component reference.
+
+  For example: "C12" => ("C", 12, None). "Cfoo" => ("C", sys.maxsize, "").
+  """
+  m = split_ref.r.match(ref)
+  if not m:
+    return (ref, sys.maxsize, "")
+  groups = m.groups()
+  ref = groups[0]
+  val = groups[1]
+  rest = groups[2]
+  try:
+    return (ref, int(val), rest)
+  except ValueError:
+    pass
+
+    return (ref, val, rest)
+split_ref.r = re.compile("([A-Za-z]+)([0-9]+)(.*)")
+
 class Part:
   def __init__(self, name):
     self.name = name
@@ -22,19 +46,26 @@ class Library:
       self.parts[part] = p
       return p
 
-class Bom:
+class Bom(object):
   def __init__(self):
     self.libraries = {}
-    self.components = {}
+    self._components = {}
 
   def add_component(self, component):
-    self.components[component.ref] = component
+    self._components[component.ref] = component
 
   def get_component(self, name):
-    return self.components[name]
+    return self._components[name]
 
   def get_components(self):
-    return self.components
+    return self._components
+
+  def all_field_names(self):
+    fields = set(['ref', 'value'])
+    for c in self._components.values():
+      for f in c.fields:
+        fields.add(f)
+    return fields
 
   def find_library(self, name):
     try:
@@ -44,6 +75,27 @@ class Bom:
       self.libraries[name] = lib
       return lib
 
+  def to_pandas(self, ref_field_name = None, value_field_name = None):
+    import pandas
+
+    ref_field_name = ref_field_name or "ref"
+    value_field_name = value_field_name or "value"
+
+    fields = self.all_field_names()
+    data = {k: [] for k in fields}
+    refs = []
+    values = []
+    for ref, c in self.get_components().items():
+      refs.append(c.ref)
+      values.append(c.value)
+      for field in fields:
+        data[field].append(c[field] if field in c else None)
+
+#    del data[ref_field_name]
+    data[ref_field_name] = refs
+    data[value_field_name] = values
+    return pandas.DataFrame(data=data, index=refs)
+
 class Comp:
   def __init__(self, ref, value, library, part, footprint):
     self.ref = ref
@@ -56,3 +108,18 @@ class Comp:
   def add_field(self, key, value):
     self.fields[key] = value
 
+  def __contains__(self, key):
+    if key == 'ref':
+      return self.ref is not None
+    elif key == 'value':
+      return self.value is not None
+    else:
+      return key in self.fields
+
+  def __getitem__(self, key):
+    if key == 'ref':
+      return self.ref
+    elif key == 'value':
+      return self.value
+    else:
+      return self.fields[key]