Merge pull request #53 from santisq/52-add-cancellation-sort-by-ascending-registrykeys

santisq · web-flow · commit 8644ff8133af · 2025-09-16T21:03:27.000-03:00
Add Ctrl+C Cancellation Support and Ascending Sort Order for `Get-PSTreeRegistry`
diff --git a/module/PSTree.psd1 b/module/PSTree.psd1
@@ -16,7 +16,7 @@
     }
 
     # Version number of this module.
-    ModuleVersion = '2.2.6'
+    ModuleVersion = '2.2.7'
 
     # Supported PSEditions
     # CompatiblePSEditions = @()
diff --git a/src/PSTree/Commands/GetPSTreeCommand.cs b/src/PSTree/Commands/GetPSTreeCommand.cs
@@ -14,7 +14,7 @@ public sealed class GetPSTreeCommand : TreeCommandBase
 {
     private readonly Stack<TreeDirectory> _stack = new();
 
-    private readonly Cache<TreeFileSystemInfo, TreeFile> _cache = new();
+    private readonly TreeBuilder<TreeFileSystemInfo, TreeFile> _builder = new();
 
     private readonly TreeComparer _comparer = new();
 
@@ -65,11 +65,11 @@ protected override void ProcessRecord()
 
     private TreeFileSystemInfo[] Traverse(TreeDirectory directory)
     {
-        _cache.Clear();
+        _builder.Clear();
         directory.PushToStack(_stack);
         string source = directory.FullName;
 
-        while (_stack.Count > 0)
+        while (_stack.Count > 0 && !Canceled)
         {
             TreeDirectory next = _stack.Pop();
             int level = next.Depth + 1;
@@ -111,7 +111,7 @@ private TreeFileSystemInfo[] Traverse(TreeDirectory directory)
                                 .Create(fileInfo, source, level)
                                 .AddParent<TreeFile>(next)
                                 .SetIncludeFlagIf(WithInclude)
-                                .AddToCache(_cache);
+                                .AddToCache(_builder);
                         }
 
                         continue;
@@ -138,15 +138,15 @@ private TreeFileSystemInfo[] Traverse(TreeDirectory directory)
 
                 if (next.Depth <= Depth)
                 {
-                    _cache.Add(next);
-                    _cache.Flush();
+                    _builder.Add(next);
+                    _builder.Flush();
                 }
             }
             catch (Exception exception)
             {
                 if (next.Depth <= Depth)
                 {
-                    _cache.Add(next);
+                    _builder.Add(next);
                 }
 
                 WriteError(exception.ToEnumerationError(next));
@@ -160,7 +160,7 @@ private TreeFileSystemInfo[] Traverse(TreeDirectory directory)
 
     private TreeFileSystemInfo[] GetTree(bool includeCondition)
     {
-        TreeFileSystemInfo[] result = _cache.GetResult(includeCondition);
+        TreeFileSystemInfo[] result = _builder.GetTree(includeCondition);
         return result.Format(GetItemCount(result));
     }
 
diff --git a/src/PSTree/Commands/GetPSTreeRegistryCommand.cs b/src/PSTree/Commands/GetPSTreeRegistryCommand.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using System.Management.Automation;
 using System.Security;
 using Microsoft.Win32;
@@ -18,7 +19,7 @@ namespace PSTree.Commands;
 public sealed class GetPSTreeRegistryCommand : TreeCommandBase
 {
 #if WINDOWS
-    private readonly Cache<TreeRegistryBase, TreeRegistryValue> _cache = new();
+    private readonly TreeBuilder<TreeRegistryBase, TreeRegistryValue> _builder = new();
 
     private readonly Stack<(TreeRegistryKey, RegistryKey)> _stack = [];
 #endif
@@ -55,15 +56,15 @@ protected override void ProcessRecord()
 
     private TreeRegistryBase[] Traverse(RegistryKey registryKey)
     {
-        _cache.Clear();
+        _builder.Clear();
 
         registryKey.
             CreateTreeKey(System.IO.Path.GetFileName(registryKey.Name)).
             PushToStack(_stack);
 
         string source = registryKey.Name;
 
-        while (_stack.Count > 0)
+        while (_stack.Count > 0 && !Canceled)
         {
             (TreeRegistryKey tree, registryKey) = _stack.Pop();
             int depth = tree.Depth + 1;
@@ -87,11 +88,11 @@ private TreeRegistryBase[] Traverse(RegistryKey registryKey)
                         new TreeRegistryValue(registryKey, value, source, depth)
                             .AddParent<TreeRegistryValue>(tree)
                             .SetIncludeFlagIf(WithInclude)
-                            .AddToCache(_cache);
+                            .AddToCache(_builder);
                     }
 
                 PushKeys:
-                    foreach (string keyname in registryKey.GetSubKeyNames())
+                    foreach (string keyname in registryKey.EnumerateKeys())
                     {
                         if (ShouldExclude(keyname))
                         {
@@ -121,11 +122,11 @@ private TreeRegistryBase[] Traverse(RegistryKey registryKey)
                 }
             }
 
-            _cache.Add(tree);
-            _cache.Flush();
+            _builder.Add(tree);
+            _builder.Flush();
         }
 
-        return _cache.GetResult(WithInclude && !KeysOnly).Format();
+        return _builder.GetTree(WithInclude && !KeysOnly).Format();
     }
 
     private bool ShouldSkipValue(string value) => ShouldExclude(value) || !ShouldInclude(value);
diff --git a/src/PSTree/Extensions/TreeExtensions.cs b/src/PSTree/Extensions/TreeExtensions.cs
@@ -1,9 +1,11 @@
 ﻿using System;
 using System.Collections.Generic;
-using System.Text;
 using Microsoft.Win32;
+using System.Linq;
 #if NETCOREAPP
 using System.Runtime.CompilerServices;
+#else
+using System.Text;
 #endif
 
 namespace PSTree.Extensions;
@@ -86,7 +88,7 @@ internal static TreeFileSystemInfo[] Format(
         return tree;
     }
 
-    internal static void AddToCache<TBase, TLeaf>(this TLeaf leaf, Cache<TBase, TLeaf> cache)
+    internal static void AddToCache<TBase, TLeaf>(this TLeaf leaf, TreeBuilder<TBase, TLeaf> cache)
         where TLeaf : TBase
         where TBase : ITree
     {
@@ -186,6 +188,13 @@ internal static TreeRegistryBase[] Format(
         return tree;
     }
 
+    internal static IEnumerable<string> EnumerateKeys(this RegistryKey registryKey) =>
+#if NET6_0_OR_GREATER
+        registryKey.GetSubKeyNames().OrderDescending();
+#else
+        registryKey.GetSubKeyNames().OrderByDescending(e => e);
+#endif
+
     internal static (TreeRegistryKey, RegistryKey) CreateTreeKey(
         this RegistryKey key, string name) =>
         (new TreeRegistryKey(key, name, key.Name), key);
diff --git a/src/PSTree/TreeBuilder.cs b/src/PSTree/TreeBuilder.cs
@@ -3,7 +3,7 @@
 
 namespace PSTree;
 
-internal sealed class Cache<TBase, TLeaf>
+internal sealed class TreeBuilder<TBase, TLeaf>
     where TLeaf : TBase
     where TBase : ITree
 {
@@ -15,7 +15,7 @@ internal sealed class Cache<TBase, TLeaf>
 
     internal void Add(TBase container) => _items.Add(container);
 
-    internal TBase[] GetResult(bool includeCondition) =>
+    internal TBase[] GetTree(bool includeCondition) =>
         includeCondition
             ? [.. _items.Where(static e => e.Include)]
             : [.. _items];
diff --git a/src/PSTree/TreeCommandBase.cs b/src/PSTree/TreeCommandBase.cs
@@ -15,12 +15,12 @@ public abstract class TreeCommandBase : PSCmdlet
 
     private WildcardPattern[]? _includePatterns;
 
+    private string[]? _paths;
+
     protected const string PathSet = "Path";
 
     protected const string LiteralPathSet = "LiteralPath";
 
-    protected string[]? _paths;
-
     protected bool WithExclude { get; private set; }
 
     protected bool WithInclude { get; private set; }
@@ -30,6 +30,8 @@ protected bool IsLiteral
         get => MyInvocation.BoundParameters.ContainsKey(nameof(LiteralPath));
     }
 
+    protected bool Canceled { get; set; }
+
     [Parameter(
         ParameterSetName = PathSet,
         Position = 0,
@@ -98,6 +100,8 @@ protected override void BeginProcessing()
         }
     }
 
+    protected override void StopProcessing() => Canceled = true;
+
     protected IEnumerable<(ProviderInfo, string)> EnumerateResolvedPaths()
     {
         Collection<string> resolvedPaths;
diff --git a/tests/GetPSTreeCommand.tests.ps1 b/tests/GetPSTreeCommand.tests.ps1
@@ -1,10 +1,13 @@
-﻿$ErrorActionPreference = 'Stop'
+﻿using namespace System.IO
+using namespace System.Linq
 
-$moduleName = (Get-Item ([IO.Path]::Combine($PSScriptRoot, '..', 'module', '*.psd1'))).BaseName
-$manifestPath = [IO.Path]::Combine($PSScriptRoot, '..', 'output', $moduleName)
+$ErrorActionPreference = 'Stop'
+
+$moduleName = (Get-Item ([Path]::Combine($PSScriptRoot, '..', 'module', '*.psd1'))).BaseName
+$manifestPath = [Path]::Combine($PSScriptRoot, '..', 'output', $moduleName)
 
 Import-Module $manifestPath
-Import-Module ([System.IO.Path]::Combine($PSScriptRoot, 'shared.psm1'))
+Import-Module ([Path]::Combine($PSScriptRoot, 'shared.psm1'))
 
 Describe 'Get-PSTree' {
     BeforeAll {
@@ -33,7 +36,7 @@ Describe 'Get-PSTree' {
             $newFolder
             New-Item @fileSplat
         } | ForEach-Object {
-            $_.Attributes = $_.Attributes -bor [System.IO.FileAttributes]::Hidden
+            $_.Attributes = $_.Attributes -bor [FileAttributes]::Hidden
         }
     }
 
@@ -109,9 +112,9 @@ Describe 'Get-PSTree' {
         $exclude = '*tools*', '*build*', '*.ps1'
         Get-PSTree $testPath -Exclude * | Should -HaveCount 1
         Get-PSTree $testPath -Exclude $exclude -Recurse | ForEach-Object {
-            [System.Linq.Enumerable]::Any(
+            [Enumerable]::Any(
                 [string[]] $exclude,
-                [System.Func[string, bool]] { $_.Name -like $args[0] })
+                [Func[string, bool]] { $_.Name -like $args[0] })
         } | Should -Not -BeTrue
 
         Get-ChildItem $testPath -Filter *.ps1 -Recurse |
@@ -122,9 +125,9 @@ Describe 'Get-PSTree' {
     It 'Includes child items with -Include parameter' {
         $include = '*.ps1', '*.cs'
         Get-PSTree $testPath -Include $include -Recurse | ForEach-Object {
-            [System.Linq.Enumerable]::Any(
+            [Enumerable]::Any(
                 [string[]] $include,
-                [System.Func[string, bool]] {
+                [Func[string, bool]] {
                     $_.Name -like $args[0] -or $_ -is [PSTree.TreeDirectory]
                 }
             )
@@ -176,4 +179,23 @@ Describe 'Get-PSTree' {
         $testHiddenFolder | Get-PSTree -Recurse -Force |
             Should -HaveCount 21
     }
+
+    It 'Should be able to Cancel the cmdlet' {
+        Measure-Command {
+            $ps = [powershell]::Create().AddScript({
+                Import-Module $args[0]
+
+                $roots = Get-PSDrive |
+                    Where-Object { $_.Provider.Name -eq 'FileSystem' } |
+                    ForEach-Object Root
+
+                Get-PSTree $roots -Recurse -ErrorAction SilentlyContinue
+            }).AddArgument($manifestPath)
+
+            $task = $ps.BeginInvoke()
+            Start-Sleep 0.5
+            $ps.Stop()
+            try { $ps.EndInvoke($task) } catch { }
+        } | Should -BeLessThan ([timespan] '00:00:01')
+    }
 }
diff --git a/tests/GetPSTreeRegistryCommand.tests.ps1 b/tests/GetPSTreeRegistryCommand.tests.ps1
diff --git a/tools/requiredModules.psd1 b/tools/requiredModules.psd1

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@`
`16`	`16`	`}`
`17`	`17`
`18`	`18`	`# Version number of this module.`
`19`		`- ModuleVersion = '2.2.6'`
	`19`	`+ ModuleVersion = '2.2.7'`
`20`	`20`
`21`	`21`	`# Supported PSEditions`
`22`	`22`	`# CompatiblePSEditions = @()`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ public sealed class GetPSTreeCommand : TreeCommandBase`
`14`	`14`	`{`
`15`	`15`	`private readonly Stack<TreeDirectory> _stack = new();`
`16`	`16`
`17`		`- private readonly Cache<TreeFileSystemInfo, TreeFile> _cache = new();`
	`17`	`+ private readonly TreeBuilder<TreeFileSystemInfo, TreeFile> _builder = new();`
`18`	`18`
`19`	`19`	`private readonly TreeComparer _comparer = new();`
`20`	`20`
`@@ -65,11 +65,11 @@ protected override void ProcessRecord()`
`65`	`65`
`66`	`66`	`private TreeFileSystemInfo[] Traverse(TreeDirectory directory)`
`67`	`67`	`{`
`68`		`- _cache.Clear();`
	`68`	`+ _builder.Clear();`
`69`	`69`	`directory.PushToStack(_stack);`
`70`	`70`	`string source = directory.FullName;`
`71`	`71`
`72`		`- while (_stack.Count > 0)`
	`72`	`+ while (_stack.Count > 0 && !Canceled)`
`73`	`73`	`{`
`74`	`74`	`TreeDirectory next = _stack.Pop();`
`75`	`75`	`int level = next.Depth + 1;`
`@@ -111,7 +111,7 @@ private TreeFileSystemInfo[] Traverse(TreeDirectory directory)`
`111`	`111`	`.Create(fileInfo, source, level)`
`112`	`112`	`.AddParent<TreeFile>(next)`
`113`	`113`	`.SetIncludeFlagIf(WithInclude)`
`114`		`- .AddToCache(_cache);`
	`114`	`+ .AddToCache(_builder);`
`115`	`115`	`}`
`116`	`116`
`117`	`117`	`continue;`
`@@ -138,15 +138,15 @@ private TreeFileSystemInfo[] Traverse(TreeDirectory directory)`
`138`	`138`
`139`	`139`	`if (next.Depth <= Depth)`
`140`	`140`	`{`
`141`		`- _cache.Add(next);`
`142`		`- _cache.Flush();`
	`141`	`+ _builder.Add(next);`
	`142`	`+ _builder.Flush();`
`143`	`143`	`}`
`144`	`144`	`}`
`145`	`145`	`catch (Exception exception)`
`146`	`146`	`{`
`147`	`147`	`if (next.Depth <= Depth)`
`148`	`148`	`{`
`149`		`- _cache.Add(next);`
	`149`	`+ _builder.Add(next);`
`150`	`150`	`}`
`151`	`151`
`152`	`152`	`WriteError(exception.ToEnumerationError(next));`
`@@ -160,7 +160,7 @@ private TreeFileSystemInfo[] Traverse(TreeDirectory directory)`
`160`	`160`
`161`	`161`	`private TreeFileSystemInfo[] GetTree(bool includeCondition)`
`162`	`162`	`{`
`163`		`- TreeFileSystemInfo[] result = _cache.GetResult(includeCondition);`
	`163`	`+ TreeFileSystemInfo[] result = _builder.GetTree(includeCondition);`
`164`	`164`	`return result.Format(GetItemCount(result));`
`165`	`165`	`}`
`166`	`166`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,7 @@`
`1`	`1`	`using System;`
`2`	`2`	`using System.Collections.Generic;`
`3`	`3`	`using System.Diagnostics.CodeAnalysis;`
	`4`	`+using System.Linq;`
`4`	`5`	`using System.Management.Automation;`
`5`	`6`	`using System.Security;`
`6`	`7`	`using Microsoft.Win32;`
`@@ -18,7 +19,7 @@ namespace PSTree.Commands;`
`18`	`19`	`public sealed class GetPSTreeRegistryCommand : TreeCommandBase`
`19`	`20`	`{`
`20`	`21`	`#if WINDOWS`
`21`		`- private readonly Cache<TreeRegistryBase, TreeRegistryValue> _cache = new();`
	`22`	`+ private readonly TreeBuilder<TreeRegistryBase, TreeRegistryValue> _builder = new();`
`22`	`23`
`23`	`24`	`private readonly Stack<(TreeRegistryKey, RegistryKey)> _stack = [];`
`24`	`25`	`#endif`
`@@ -55,15 +56,15 @@ protected override void ProcessRecord()`
`55`	`56`
`56`	`57`	`private TreeRegistryBase[] Traverse(RegistryKey registryKey)`
`57`	`58`	`{`
`58`		`- _cache.Clear();`
	`59`	`+ _builder.Clear();`
`59`	`60`
`60`	`61`	`registryKey.`
`61`	`62`	`CreateTreeKey(System.IO.Path.GetFileName(registryKey.Name)).`
`62`	`63`	`PushToStack(_stack);`
`63`	`64`
`64`	`65`	`string source = registryKey.Name;`
`65`	`66`
`66`		`- while (_stack.Count > 0)`
	`67`	`+ while (_stack.Count > 0 && !Canceled)`
`67`	`68`	`{`
`68`	`69`	`(TreeRegistryKey tree, registryKey) = _stack.Pop();`
`69`	`70`	`int depth = tree.Depth + 1;`
`@@ -87,11 +88,11 @@ private TreeRegistryBase[] Traverse(RegistryKey registryKey)`
`87`	`88`	`new TreeRegistryValue(registryKey, value, source, depth)`
`88`	`89`	`.AddParent<TreeRegistryValue>(tree)`
`89`	`90`	`.SetIncludeFlagIf(WithInclude)`
`90`		`- .AddToCache(_cache);`
	`91`	`+ .AddToCache(_builder);`
`91`	`92`	`}`
`92`	`93`
`93`	`94`	`PushKeys:`
`94`		`- foreach (string keyname in registryKey.GetSubKeyNames())`
	`95`	`+ foreach (string keyname in registryKey.EnumerateKeys())`
`95`	`96`	`{`
`96`	`97`	`if (ShouldExclude(keyname))`
`97`	`98`	`{`
`@@ -121,11 +122,11 @@ private TreeRegistryBase[] Traverse(RegistryKey registryKey)`
`121`	`122`	`}`
`122`	`123`	`}`
`123`	`124`
`124`		`- _cache.Add(tree);`
`125`		`- _cache.Flush();`
	`125`	`+ _builder.Add(tree);`
	`126`	`+ _builder.Flush();`
`126`	`127`	`}`
`127`	`128`
`128`		`- return _cache.GetResult(WithInclude && !KeysOnly).Format();`
	`129`	`+ return _builder.GetTree(WithInclude && !KeysOnly).Format();`
`129`	`130`	`}`
`130`	`131`
`131`	`132`	`private bool ShouldSkipValue(string value) => ShouldExclude(value) \|\| !ShouldInclude(value);`
Original file line number	Diff line number	Diff line change
`@@ -15,12 +15,12 @@ public abstract class TreeCommandBase : PSCmdlet`
`15`	`15`
`16`	`16`	`private WildcardPattern[]? _includePatterns;`
`17`	`17`
	`18`	`+ private string[]? _paths;`
	`19`	`+`
`18`	`20`	`protected const string PathSet = "Path";`
`19`	`21`
`20`	`22`	`protected const string LiteralPathSet = "LiteralPath";`
`21`	`23`
`22`		`- protected string[]? _paths;`
`23`		`-`
`24`	`24`	`protected bool WithExclude { get; private set; }`
`25`	`25`
`26`	`26`	`protected bool WithInclude { get; private set; }`
`@@ -30,6 +30,8 @@ protected bool IsLiteral`
`30`	`30`	`get => MyInvocation.BoundParameters.ContainsKey(nameof(LiteralPath));`
`31`	`31`	`}`
`32`	`32`
	`33`	`+ protected bool Canceled { get; set; }`
	`34`	`+`
`33`	`35`	`[Parameter(`
`34`	`36`	`ParameterSetName = PathSet,`
`35`	`37`	`Position = 0,`
`@@ -98,6 +100,8 @@ protected override void BeginProcessing()`
`98`	`100`	`}`
`99`	`101`	`}`
`100`	`102`
	`103`	`+ protected override void StopProcessing() => Canceled = true;`
	`104`	`+`
`101`	`105`	`protected IEnumerable<(ProviderInfo, string)> EnumerateResolvedPaths()`
`102`	`106`	`{`
`103`	`107`	`Collection<string> resolvedPaths;`