fix(tree): allow tree node to have undefined child or null child (#14482)

Author: vivian-hu-zz

src/lib/tree/data-source/flat-data-source.ts

@@ -50,7 +50,8 @@ export class MatTreeFlattener<T, F> {
   constructor(public transformFunction: (node: T, level: number) => F,
               public getLevel: (node: F) => number,
               public isExpandable: (node: F) => boolean,
-              public getChildren: (node: T) => Observable<T[]> | T[]) {}
+              public getChildren: (node: T) =>
+                  Observable<T[]> | T[] | undefined | null) {}
 
   _flattenNode(node: T, level: number,
                resultNodes: F[], parentMap: boolean[]): F[] {
@@ -59,12 +60,14 @@ export class MatTreeFlattener<T, F> {
 
     if (this.isExpandable(flatNode)) {
       const childrenNodes = this.getChildren(node);
-      if (Array.isArray(childrenNodes)) {
-        this._flattenChildren(childrenNodes, level, resultNodes, parentMap);
-      } else {
-        childrenNodes.pipe(take(1)).subscribe(children => {
-          this._flattenChildren(children, level, resultNodes, parentMap);
-        });
+      if (childrenNodes) {
+        if (Array.isArray(childrenNodes)) {
+          this._flattenChildren(childrenNodes, level, resultNodes, parentMap);
+        } else {
+          childrenNodes.pipe(take(1)).subscribe(children => {
+            this._flattenChildren(children, level, resultNodes, parentMap);
+          });
+        }
       }
     }
     return resultNodes;

src/lib/tree/tree.spec.ts

@@ -216,6 +216,27 @@ describe('MatTree', () => {
     });
   });
 
+  describe('flat tree with undefined or null children', () => {
+    describe('should initialize', () => {
+      let fixture: ComponentFixture<MatTreeWithNullOrUndefinedChild >;
+
+      beforeEach(() => {
+        configureMatTreeTestingModule([MatTreeWithNullOrUndefinedChild ]);
+        fixture = TestBed.createComponent(MatTreeWithNullOrUndefinedChild );
+        treeElement = fixture.nativeElement.querySelector('mat-tree');
+
+        fixture.detectChanges();
+      });
+
+      it('with rendered dataNodes', () => {
+        const nodes = getNodes(treeElement);
+
+        expect(nodes).toBeDefined('Expect nodes to be defined');
+        expect(nodes[0].classList).toContain('customNodeClass');
+      });
+    });
+  });
+
   describe('nested tree', () => {
     describe('should initialize', () => {
       let fixture: ComponentFixture<NestedMatTreeApp>;
@@ -610,6 +631,85 @@ class SimpleMatTreeApp {
   }
 }
 
+interface FoodNode {
+  name: string;
+  children?: FoodNode[] | null;
+}
+
+/** Flat node with expandable and level information */
+interface ExampleFlatNode {
+  expandable: boolean;
+  name: string;
+  level: number;
+}
+
+/**
+ * Food data with nested structure.
+ * Each node has a name and an optiona list of children.
+ */
+const TREE_DATA: FoodNode[] = [
+  {
+    name: 'Fruit',
+    children: [
+      {name: 'Apple'},
+      {name: 'Banana'},
+      {name: 'Fruit loops',
+       children: null},
+    ]
+  }, {
+    name: 'Vegetables',
+    children: [
+      {
+        name: 'Green',
+        children: [
+          {name: 'Broccoli'},
+          {name: 'Brussel sprouts'},
+        ]
+      }, {
+        name: 'Orange',
+        children: [
+          {name: 'Pumpkins'},
+          {name: 'Carrots'},
+        ]
+      },
+    ]
+  },
+];
+
+@Component({
+  template: `
+    <mat-tree [dataSource]="dataSource" [treeControl]="treeControl">
+      <mat-tree-node *matTreeNodeDef="let node" class="customNodeClass"
+                     matTreeNodePadding matTreeNodeToggle>
+        {{node.name}}
+      </mat-tree-node>
+    </mat-tree>
+  `
+})
+class MatTreeWithNullOrUndefinedChild {
+  private transformer = (node: FoodNode, level: number) => {
+    return {
+      expandable: !!node.children,
+      name: node.name,
+      level: level,
+    };
+  }
+
+  treeControl = new FlatTreeControl<ExampleFlatNode>(
+    node => node.level, node => node.expandable);
+
+  treeFlattener = new MatTreeFlattener(
+     this.transformer, node => node.level, node => node.expandable, node => node.children);
+
+  dataSource = new MatTreeFlatDataSource(this.treeControl, this.treeFlattener, TREE_DATA);
+
+  constructor() {
+    this.dataSource.data = TREE_DATA;
+  }
+
+  hasChild = (_: number, node: ExampleFlatNode) => node.expandable;
+}
+
 @Component({
   template: `
     <mat-tree [dataSource]="dataSource" [treeControl]="treeControl">