simplifies logic to get item and totalitem count

santisq · Jan 19, 2025 · 3567b8b · 3567b8b
1 parent e04d1e9
commit 3567b8b
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 36 deletions.
diff --git a/src/PSTree/Cache.cs b/src/PSTree/Cache.cs
@@ -23,19 +23,35 @@ internal void Flush()
         }
     }
 
-    internal PSTreeFileSystemInfo[] GetTree(bool condition) =>
-        condition
-            ? _items.Where(IsIncluded).ToArray().Format()
-            : _items.ToArray().Format();
+    internal PSTreeFileSystemInfo[] GetTree(bool condition)
+    {
+        PSTreeFileSystemInfo[] result = condition
+            ? [.. _items.Where(static e => e.ShouldInclude)]
+            : [.. _items];
+
+        return result.Format(GetItemCount(result));
+    }
 
-    private static bool IsIncluded(PSTreeFileSystemInfo item)
+    private static Dictionary<string, int> GetItemCount(PSTreeFileSystemInfo[] items)
     {
-        if (item.ShouldInclude && item is PSTreeDirectory dir)
+        Dictionary<string, int> counts = [];
+        foreach (PSTreeFileSystemInfo item in items)
         {
-            dir.IncrementItemCount();
+            string? path = item.ParentNode?.FullName;
+            if (path is null)
+            {
+                continue;
+            }
+
+            if (!counts.ContainsKey(path))
+            {
+                counts[path] = 0;
+            }
+
+            counts[path]++;
         }
 
-        return item.ShouldInclude;
+        return counts;
     }
 
     internal void Clear()

diff --git a/src/PSTree/Commands/GetPSTreeCommand.cs b/src/PSTree/Commands/GetPSTreeCommand.cs
@@ -108,7 +108,6 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
             PSTreeDirectory next = _stack.Pop();
             int level = next.Depth + 1;
             long totalLength = 0;
-            int childCount = 0;
 
             try
             {
@@ -142,8 +141,6 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
 
                         if (keepProcessing)
                         {
-                            childCount++;
-
                             PSTreeFile file = PSTreeFile
                                 .Create(fileInfo, source, level)
                                 .AddParent(next)
@@ -167,14 +164,12 @@ private PSTreeFileSystemInfo[] Traverse(PSTreeDirectory directory)
                     if (keepProcessing && Directory || !_withInclude)
                     {
                         dir.ShouldInclude = true;
-                        childCount++;
                     }
 
                     _stack.Push(dir);
                 }
 
                 next.Length = totalLength;
-                next.IndexCount(childCount);
 
                 if (RecursiveSize)
                 {

diff --git a/src/PSTree/Extensions/TreeExtensions.cs b/src/PSTree/Extensions/TreeExtensions.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Text;
 
 namespace PSTree.Extensions;
@@ -19,12 +20,21 @@ internal static string Indent(this string inputString, int indentation)
             .ToString();
     }
 
-    internal static PSTreeFileSystemInfo[] Format(this PSTreeFileSystemInfo[] tree)
+    internal static PSTreeFileSystemInfo[] Format(
+        this PSTreeFileSystemInfo[] tree,
+        Dictionary<string, int> itemCounts)
     {
         int index;
         for (int i = 0; i < tree.Length; i++)
         {
             PSTreeFileSystemInfo current = tree[i];
+
+            if (current is PSTreeDirectory directory &&
+                itemCounts.TryGetValue(directory.FullName, out int count))
+            {
+                directory.IndexCount(count);
+            }
+
             if ((index = current.Hierarchy.IndexOf('└')) == -1)
             {
                 continue;

diff --git a/src/PSTree/PSTreeDirectory.cs b/src/PSTree/PSTreeDirectory.cs
@@ -58,7 +58,7 @@ internal static PSTreeDirectory Create(DirectoryInfo dir, string source, int dep
 
     internal PSTreeDirectory AddParent(PSTreeDirectory parent)
     {
-        _parent = parent;
+        ParentNode = parent;
         return this;
     }
 
@@ -67,15 +67,15 @@ internal void IndexCount(int count)
         ItemCount = count;
         TotalItemCount = count;
 
-        for (PSTreeDirectory? i = _parent; i is not null; i = i._parent)
+        for (PSTreeDirectory? i = ParentNode; i is not null; i = i.ParentNode)
         {
             i.TotalItemCount += count;
         }
     }
 
     internal void IndexLength(long length)
     {
-        for (PSTreeDirectory? i = _parent; i is not null; i = i._parent)
+        for (PSTreeDirectory? i = ParentNode; i is not null; i = i.ParentNode)
         {
             i.Length += length;
         }
@@ -85,7 +85,7 @@ internal void SetIncludeFlag()
     {
         ShouldInclude = true;
 
-        for (PSTreeDirectory? i = _parent; i is not null; i = i._parent)
+        for (PSTreeDirectory? i = ParentNode; i is not null; i = i.ParentNode)
         {
             if (i.ShouldInclude)
             {
@@ -95,19 +95,4 @@ internal void SetIncludeFlag()
             i.ShouldInclude = true;
         }
     }
-
-    internal void IncrementItemCount()
-    {
-        if (_parent is null)
-        {
-            return;
-        }
-
-        _parent.ItemCount++;
-
-        for (PSTreeDirectory? i = _parent; i is not null; i = i._parent)
-        {
-            i.TotalItemCount++;
-        }
-    }
 }
diff --git a/src/PSTree/PSTreeFile.cs b/src/PSTree/PSTreeFile.cs
@@ -40,15 +40,15 @@ internal static PSTreeFile Create(FileInfo file, string source, int depth)
 
     internal PSTreeFile AddParent(PSTreeDirectory parent)
     {
-        _parent = parent;
+        ParentNode = parent;
         return this;
     }
 
     internal PSTreeFile SetIncludeFlagIf(bool condition)
     {
         if (condition)
         {
-            _parent?.SetIncludeFlag();
+            ParentNode?.SetIncludeFlag();
         }
 
         return this;

diff --git a/src/PSTree/PSTreeFileSystemInfo.cs b/src/PSTree/PSTreeFileSystemInfo.cs
@@ -2,7 +2,7 @@ namespace PSTree;
 
 public abstract class PSTreeFileSystemInfo(string hierarchy, string source)
 {
-    protected PSTreeDirectory? _parent;
+    internal PSTreeDirectory? ParentNode { get; set; }
 
     internal bool ShouldInclude { get; set; }