almost there, need to work on item counts

santisq · Jan 17, 2025 · ffc7c85 · ffc7c85
1 parent 6588c0e
commit ffc7c85
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 25 deletions.
diff --git a/src/PSTree/Cache.cs b/src/PSTree/Cache.cs
@@ -10,11 +10,11 @@ internal sealed class Cache
 
     private readonly List<PSTreeFile> _files = [];
 
-    internal void AddFile(PSTreeFile file) => _files.Add(file);
+    internal void Add(PSTreeFile file) => _files.Add(file);
 
-    internal void Add(PSTreeFileSystemInfo item) => _items.Add(item);
+    internal void Add(PSTreeDirectory directory) => _items.Add(directory);
 
-    internal void TryAddFiles()
+    internal void Flush()
     {
         if (_files.Count > 0)
         {
@@ -23,10 +23,10 @@ internal void TryAddFiles()
         }
     }
 
-    internal PSTreeFileSystemInfo[] GetTree() => _items
-        .Where(e => e._shouldInclude)
-        .ToArray()
-        .ConvertToTree();
+    internal PSTreeFileSystemInfo[] GetTree(bool filterInclude) =>
+        filterInclude
+            ? _items.Where(e => e.ShouldInclude).ToArray().ConvertToTree()
+            : _items.ToArray().ConvertToTree();
 
     internal void Clear()
     {

diff --git a/src/PSTree/Commands/GetPSTreeCommand.cs b/src/PSTree/Commands/GetPSTreeCommand.cs
@@ -76,7 +76,12 @@ protected override void ProcessRecord()
         {
             if (File.Exists(path))
             {
-                WriteObject(PSTreeFile.Create(path));
+                FileInfo file = new(path);
+                if (!ShouldExclude(file) && ShouldInclude(file))
+                {
+                    WriteObject(PSTreeFile.Create(file, path));
+                }
+
                 continue;
             }
 
@@ -118,6 +123,7 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
                     {
                         if (Directory.IsPresent)
                         {
+                            size += fileInfo.Length;
                             continue;
                         }
 
@@ -128,9 +134,10 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
 
                             PSTreeFile file = PSTreeFile
                                 .Create(fileInfo, source, level)
-                                .WithParent(next);
+                                .WithParent(next)
+                                .WithIncludeFlagIf(_includePatterns is not null);
 
-                            _cache.AddFile(file);
+                            _cache.Add(file);
                         }
 
                         continue;
@@ -144,7 +151,7 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
 
                         if (_includePatterns is null)
                         {
-                            dir._shouldInclude = true;
+                            dir.ShouldInclude = true;
                             childCount++;
                         }
 
@@ -153,7 +160,7 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
                 }
 
                 next.Length = size;
-                next.IndexItemCount(childCount);
+                next.IndexCount(childCount);
 
                 if (RecursiveSize.IsPresent)
                 {
@@ -163,7 +170,7 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
                 if (next.Depth <= Depth)
                 {
                     _cache.Add(next);
-                    _cache.TryAddFiles();
+                    _cache.Flush();
                 }
             }
             catch (Exception exception)
@@ -177,7 +184,7 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
             }
         }
 
-        return _cache.GetTree();
+        return _cache.GetTree(_includePatterns is not null);
     }
 
     private static bool MatchAny(string name, WildcardPattern[] patterns)

diff --git a/src/PSTree/Extensions/TreeExtensions.cs b/src/PSTree/Extensions/TreeExtensions.cs
@@ -21,6 +21,7 @@ internal static PSTreeFileSystemInfo[] ConvertToTree(
     {
         int index;
         PSTreeFileSystemInfo current;
+
         for (int i = 0; i < inputObject.Length; i++)
         {
             current = inputObject[i];

diff --git a/src/PSTree/PSTreeDirectory.cs b/src/PSTree/PSTreeDirectory.cs
@@ -22,7 +22,9 @@ private PSTreeDirectory(
     private PSTreeDirectory(
         DirectoryInfo dir, string hierarchy, string source)
         : base(dir, hierarchy, source)
-    { }
+    {
+        ShouldInclude = true;
+    }
 
     public IEnumerable<FileInfo> EnumerateFiles() =>
         Instance.EnumerateFiles();
@@ -59,7 +61,7 @@ internal PSTreeDirectory WithParent(PSTreeDirectory parent)
         return this;
     }
 
-    internal void IndexItemCount(int count)
+    internal void IndexCount(int count)
     {
         ItemCount = count;
         TotalItemCount = count;
@@ -80,15 +82,15 @@ internal void IndexLength(long length)
 
     internal void SetIncludeFlag()
     {
-        _shouldInclude = true;
+        ShouldInclude = true;
         for (PSTreeDirectory? parent = _parent; parent is not null; parent = parent._parent)
         {
-            if (parent._shouldInclude)
+            if (parent.ShouldInclude)
             {
                 return;
             }
 
-            parent._shouldInclude = true;
+            parent.ShouldInclude = true;
         }
     }
 }
diff --git a/src/PSTree/PSTreeFile.cs b/src/PSTree/PSTreeFile.cs
@@ -15,19 +15,17 @@ private PSTreeFile(
         : base(file, hierarchy, source)
     {
         Length = file.Length;
-        _shouldInclude = true;
+        ShouldInclude = true;
     }
 
     private PSTreeFile(
         FileInfo file, string hierarchy, string source, int depth)
         : base(file, hierarchy, source, depth)
     {
         Length = file.Length;
-        _shouldInclude = true;
+        ShouldInclude = true;
     }
 
-    internal static PSTreeFile Create(string path) => Create(new FileInfo(path), path);
-
     internal static PSTreeFile Create(FileInfo file, string source)
     {
         string styled = TreeStyle.Instance.GetColoredName(file);
@@ -43,7 +41,16 @@ internal static PSTreeFile Create(FileInfo file, string source, int depth)
     internal PSTreeFile WithParent(PSTreeDirectory parent)
     {
         _parent = parent;
-        _parent.SetIncludeFlag();
+        return this;
+    }
+
+    internal PSTreeFile WithIncludeFlagIf(bool condition)
+    {
+        if (condition)
+        {
+            _parent?.SetIncludeFlag();
+        }
+
         return this;
     }
 }
diff --git a/src/PSTree/PSTreeFileSystemInfo.cs b/src/PSTree/PSTreeFileSystemInfo.cs
@@ -4,7 +4,7 @@ public abstract class PSTreeFileSystemInfo(string hierarchy, string source)
 {
     protected PSTreeDirectory? _parent;
 
-    internal bool _shouldInclude;
+    internal bool ShouldInclude { get; set; }
 
     internal string Source { get; set; } = source;
 

diff --git a/tests/GetPSTreeCommand.tests.ps1 b/tests/GetPSTreeCommand.tests.ps1
@@ -113,6 +113,10 @@ Describe 'Get-PSTree' {
                 [string[]] $exclude,
                 [System.Func[string, bool]] { $_.FullName -like $args[0] })
         } | Should -Not -BeTrue
+
+        Get-ChildItem $testPath -Filter *.ps1 -Recurse |
+            Get-PSTree -Exclude *.ps1 |
+            Should -BeNullOrEmpty
     }
 
     It 'Includes child items with -Include parameter' {
@@ -125,6 +129,10 @@ Describe 'Get-PSTree' {
                 }
             )
         } | Should -BeTrue
+
+        Get-ChildItem $testPath -Filter *.ps1 -Recurse |
+            Get-PSTree -Include *.ps1 |
+            Should -Not -BeNullOrEmpty
     }
 
     It 'Should prioritize -Depth if used together with -Recurse' {