fix(@angular-devkit/schematics): support delegating trees in filter and partition

packages/angular_devkit/schematics/src/rules/base.ts

@@ -9,6 +9,7 @@
 import { concat, lastValueFrom, map, mapTo, toArray } from 'rxjs';
 import { FileOperator, Rule, Source } from '../engine/interface';
 import { SchematicsException } from '../exception/exception';
+import { DelegateTree } from '../tree/delegate';
 import { FilterHostTree, HostTree } from '../tree/host-tree';
 import { FileEntry, FilePredicate, MergeStrategy, Tree } from '../tree/interface';
 import { ScopedTree } from '../tree/scoped';
@@ -74,6 +75,10 @@ export function noop(): Rule {
 
 export function filter(predicate: FilePredicate<boolean>): Rule {
   return (tree: Tree) => {
+    if (tree instanceof DelegateTree) {
+      tree = tree.branch();
+    }
+
     if (HostTree.isHostTree(tree)) {
       return new FilterHostTree(tree, predicate);
     } else {

packages/angular_devkit/schematics/src/rules/base_spec.ts

@@ -8,11 +8,12 @@
 
 import { Path, virtualFs } from '@angular-devkit/core';
 import { lastValueFrom, of as observableOf } from 'rxjs';
+import { UnitTestTree } from '../../testing';
 import { Rule, SchematicContext, Source } from '../engine/interface';
 import { HostTree } from '../tree/host-tree';
 import { MergeStrategy, Tree } from '../tree/interface';
 import { empty } from '../tree/static';
-import { apply, applyToSubtree, chain, partitionApplyMerge } from './base';
+import { apply, applyToSubtree, chain, filter, partitionApplyMerge } from './base';
 import { callRule, callSource } from './call';
 import { move } from './move';
 
@@ -225,6 +226,32 @@ describe('partitionApplyMerge', () => {
     expect(result.exists('/test1')).toBe(false);
     expect(result.exists('/test2')).toBe(false);
   });
+
+  it('works with a UnitTestTree', async () => {
+    const tree = new HostTree();
+    tree.create('/test1', '');
+    tree.create('/test2', '');
+    const predicate = (path: Path) => path.indexOf('1') != -1;
+
+    const ruleYes: Rule = (tree: Tree) => {
+      expect(tree.exists('/test1')).toBe(true);
+      expect(tree.exists('/test2')).toBe(false);
+    };
+    const ruleNo: Rule = (tree: Tree) => {
+      expect(tree.exists('/test1')).toBe(false);
+      expect(tree.exists('/test2')).toBe(true);
+    };
+
+    const result = await lastValueFrom(
+      callRule(
+        partitionApplyMerge(predicate, ruleYes, ruleNo),
+        observableOf(new UnitTestTree(new UnitTestTree(tree))),
+        context,
+      ),
+    );
+    expect(result.exists('/test1')).toBe(true);
+    expect(result.exists('/test2')).toBe(true);
+  });
 });
 
 describe('applyToSubtree', () => {
@@ -242,3 +269,37 @@ describe('applyToSubtree', () => {
     expect(result.exists('a/c/file3')).toBe(true);
   });
 });
+
+describe('filter', () => {
+  it('works with a HostTree', async () => {
+    const tree = new HostTree();
+    tree.create('/file1.ts', '');
+    tree.create('/file2.js', '');
+
+    const result = await lastValueFrom(
+      callRule(
+        filter((path) => path.endsWith('.ts')),
+        observableOf(tree),
+        context,
+      ),
+    );
+    expect(result.exists('/file1.ts')).toBe(true);
+    expect(result.exists('/file2.js')).toBe(false);
+  });
+
+  it('works with a UnitTestTree', async () => {
+    const tree = new HostTree();
+    tree.create('/file1.ts', '');
+    tree.create('/file2.js', '');
+
+    const result = await lastValueFrom(
+      callRule(
+        filter((path) => path.endsWith('.ts')),
+        observableOf(new UnitTestTree(new UnitTestTree(tree))),
+        context,
+      ),
+    );
+    expect(result.exists('/file1.ts')).toBe(true);
+    expect(result.exists('/file2.js')).toBe(false);
+  });
+});

packages/angular_devkit/schematics/src/tree/static.ts

@@ -7,6 +7,7 @@
  */
 
 import { SchematicsException } from '../exception/exception';
+import { DelegateTree } from './delegate';
 import { FilterHostTree, HostTree } from './host-tree';
 import { FilePredicate, MergeStrategy, Tree } from './interface';
 
@@ -29,6 +30,10 @@ export function merge(
 }
 
 export function partition(tree: Tree, predicate: FilePredicate<boolean>): [Tree, Tree] {
+  if (tree instanceof DelegateTree) {
+    tree = tree.branch();
+  }
+
   if (tree instanceof HostTree) {
     return [
       new FilterHostTree(tree, predicate),