Create CollectionWithCategories

src/abstract/CollectionWithCategories.ts

@@ -0,0 +1,31 @@
+import { isString } from 'underscore';
+import { Collection, Model } from '../common';
+import Categories from './ModuleCategories';
+import Category, { CategoryProperties } from './ModuleCategory';
+import { isObject } from '../utils/mixins';
+
+interface ModelWithCategoryProps {
+  category?: string | CategoryProperties;
+}
+
+const CATEGORY_KEY = 'category';
+
+export abstract class CollectionWithCategories<T extends Model<ModelWithCategoryProps>> extends Collection<T> {
+  abstract get categories(): Categories;
+
+  initCategory(model: T) {
+    let category = model.get(CATEGORY_KEY);
+
+    // Ensure the category exists and it's not already initialized
+    if (category && !(category instanceof Category)) {
+      if (isString(category)) {
+        category = { id: category, label: category };
+      } else if (isObject(category) && !category.id) {
+        category.id = category.label;
+      }
+
+      const catModel = this.categories.add(category);
+      model.set(CATEGORY_KEY, catModel as any, { silent: true });
+    }
+  }
+}

src/trait_manager/model/Traits.ts

@@ -1,11 +1,12 @@
 import { isArray } from 'underscore';
-import { AddOptions, Collection } from '../../common';
+import { CollectionWithCategories } from '../../abstract/CollectionWithCategories';
+import { AddOptions } from '../../common';
 import Component from '../../dom_components/model/Component';
 import EditorModel from '../../editor/model/Editor';
 import Trait, { TraitProperties } from './Trait';
 import TraitFactory from './TraitFactory';
 
-export default class Traits extends Collection<Trait> {
+export default class Traits extends CollectionWithCategories<Trait> {
   em: EditorModel;
   target!: Component;
   tf: TraitFactory;
@@ -24,6 +25,10 @@ export default class Traits extends Collection<Trait> {
     return this.em.Traits;
   }
 
+  get categories() {
+    return this.module.categories;
+  }
+
   handleReset(coll: TraitProperties[], { previousModels = [] }: { previousModels?: Trait[] } = {}) {
     previousModels.forEach(model => model.trigger('remove'));
   }
@@ -35,6 +40,8 @@ export default class Traits extends Collection<Trait> {
     if (target) {
       model.target = target;
     }
+
+    this.initCategory(model);
   }
 
   setTarget(target: Component) {