1 files changed, 60 insertions, 1 deletions
diff --git a/src/ee/part/__init__.py b/src/ee/part/__init__.py
index 975dc68..0845b6c 100644
--- a/src/ee/part/__init__.py
+++ b/src/ee/part/__init__.py
@@ -1,3 +1,4 @@
+from functools import total_ordering
 from pathlib import Path
 from typing import List, Optional, Iterator, Union
 
@@ -8,6 +9,7 @@ from ee.xml import types
 __all__ = [
     "FactType",
     "EeValueFactType",
+    "Category",
     "Part",
     "PartDb",
     "load_db",
@@ -179,6 +181,34 @@ class ReferenceList(object):
         self.description_references.append(description)
 
 
+class Links:
+    def __init__(self, links: Optional[types.LinkList]):
+        self.links = links.linkProp if links else []
+
+    def by_rel(self, rel) -> Optional[types.Link]:
+        for l in self.links:
+            if l.relationProp == rel:
+                return l
+
+
+@total_ordering
+class Category:
+    def __init__(self, uri: str, name: str, href: Optional[str], parent: Optional["Category"]):
+        self.uri = uri
+        self.name = name
+        self.href = href
+        self.parent = parent
+
+    def __eq__(self, other):
+        return self.uri == other.uri
+
+    def __hash__(self):
+        return hash(self.uri)
+
+    def __lt__(self, other):
+        return self.uri < other.uri
+
+
 # TODO: Replace self.xml.referencesProp with ReferenceList
 class Part(object):
     def __init__(self, xml: types.Part):
@@ -412,6 +442,7 @@ class Assembly(object):
 class PartDb(object):
     def __init__(self):
         self.parts: List[Entry] = []
+        self.categories: List[Category] = []
         self.new_entries = 0
         self._assembly: Optional[Assembly] = None
 
@@ -457,6 +488,25 @@ def load_db(path: Path) -> PartDb:
     for p in part_db.partsProp.part:
         db.add_entry(p, False)
 
+    part_db.categoriesProp = part_db.categoriesProp or types.CategoryList()
+    categories_by_uri = {}
+    categories_by_parent = []
+    xml: types.Category
+    for xml in part_db.categoriesProp.category:
+        links = Links(xml.linksProp)
+
+        c = Category(xml.uri, xml.name, None, None)
+        categories_by_uri[c.uri] = c
+
+        parent = links.by_rel("parent")
+        if parent:
+            categories_by_parent.append((c, parent.urlProp))
+
+        db.categories.append(c)
+
+    for c, parent_uri in categories_by_parent:
+        c.parent = categories_by_uri[parent_uri]
+
     return db
 
 
@@ -475,8 +525,17 @@ def save_db(path: Path, db: PartDb, sort=False):
             p.clean_xml()
             part_db.parts.partProp.append(p.underlying)
 
-    if db.has_assembly:
+    if len(db.categories):
+        cs = types.CategoryList()
+        for c in db.categories:
+            links = None
+            if c.parent:
+                links = types.LinkList()
+                links.add_link(types.Link(url=c.parent.uri, relation="parent"))
+            cs.category.append(types.Category(uri=c.uri, name=c.name, links=links))
+        part_db.set_categories(cs)
 
+    if db.has_assembly:
         def to_xml(ap: AssemblyPart):
             xml = types.AssemblyPart()
             if ap.count != 0: