Folder Tree WPF Control: Part 3 - Enumeration

FolderTreeSelection-1.cs

public class FolderTreeSelection : ViewModelBase
{
    private Dictionary<Item, long> itemCounts = new Dictionary<Item, long>(); 

    public long FolderCount {
        get
        {
            return itemCounts.Sum(x => x.Value);
        }
    }

    public IEnumerable<DirectoryInfo> SelectedFolders
    {
        get
        {
            if (items.Count() > 0)
            {
                foreach (var dir in items.Where(x => x.Value.Include).Select(x => x.Value))
                {
                    yield return dir.TreeItem.DirectoryInfo;
                    itemCounts[dir] = 1;
                    OnPropertyChanged("FolderCount");
                    foreach (var subDir in EnumerateSelectedSubFolders(dir.TreeItem.DirectoryInfo))
                    {
                        yield return subDir;
                        itemCounts[dir]++;
                        OnPropertyChanged("FolderCount");
                    }
                }
            }
        }
    }

    
    internal void Add(FolderTreeItem folderTreeItem, bool include = true)
    {
        if (folderTreeItem.SelectionItem == null && include)
        {
            // node is checked for the first time, add a selection item
            folderTreeItem.SelectionItem = new Item() { TreeItem = folderTreeItem, Include = include };
            items.Add(folderTreeItem.Path, folderTreeItem.SelectionItem);

            OnChanged();
        }
        else if (folderTreeItem.SelectionItem != null && folderTreeItem.SelectionItem.Include != include)
        {
            // node with a current state is checked/unchecked
            if (!include && (folderTreeItem.Parent == null || folderTreeItem.Parent.SelectionItem == null))
            {
                // special behavior if root node is unchecked
                items.Remove(folderTreeItem.Path);
                itemCounts.Remove(folderTreeItem.SelectionItem);
                OnPropertyChanged("FolderCount");
                folderTreeItem.SelectionItem = null;
            }
            else
            {
                // determine if this node itself was specifically checked/unchecked
                var match = items.Where(x => x.Key.Equals(folderTreeItem.Path)).SingleOrDefault().Value;
                if (match != null)
                {
                    // if so, adjust it's state
                    match.Include = include;
                    if (!match.Include)
                    {
                        itemCounts.Remove(match);
                        OnPropertyChanged("FolderCount");
                    }
                }
                else
                {
                    // if not, add a new selection item
                    folderTreeItem.SelectionItem = new Item() { TreeItem = folderTreeItem, Include = include };
                    items.Add(folderTreeItem.Path, folderTreeItem.SelectionItem);
                }
            }
            OnChanged();
        }

        // process children nodes if any are loaded
        if (folderTreeItem.FoldersLoaded)
        {
            foreach (var childItem in folderTreeItem.Folders)
            {
                // all children node states should be wiped out and given a reference to this nodes state
                if (items.Remove(childItem.Path))
                {
                    OnChanged();
                }
                childItem.SelectionItem = folderTreeItem.SelectionItem;
                childItem.IsChecked = include;

                // recurse to go as many levels deep as we can
                // recursive calls will skip the code above because the selection item is already set here
                Add(childItem, include);
            }
        }

        // give the parent nodes either checked or indeterminate states
        ProcessParents(folderTreeItem);
    }
}

FolderTreeSelection.cs

public class FolderTreeSelection : ViewModelBase
{
    public IEnumerable<DirectoryInfo> SelectedFolders
    {
        get
        {
            if (items.Count() > 0)
            {
                foreach (var dir in items.Where(x => x.Value.Include).Select(x => x.Value))
                {
                    yield return dir.TreeItem.DirectoryInfo;
                    foreach (var subDir in EnumerateSelectedSubFolders(dir.TreeItem.DirectoryInfo))
                    {
                        yield return subDir;
                    }
                }
            }
        }
    }

    private IEnumerable<DirectoryInfo> EnumerateSelectedSubFolders(DirectoryInfo directoryInfo)
    {
        IEnumerable<DirectoryInfo> subDirs;
        try
        {
            subDirs = directoryInfo.EnumerateDirectories();
        }
        catch (UnauthorizedAccessException e)
        {
            subDirs = Enumerable.Empty<DirectoryInfo>();
        }
        catch (IOException e)
        {
            subDirs = Enumerable.Empty<DirectoryInfo>();
        }

        foreach (var subDir in subDirs)
        {
            if (items.ContainsKey(subDir.FullName) && !items[subDir.FullName].Include)
            {
                continue;
            }
            yield return subDir;
            foreach (var recursion in EnumerateSelectedSubFolders(subDir))
            {
                yield return recursion;
            }
        }
    }

    internal void Add(FolderTreeItem folderTreeItem, bool include = true)
    {
        if (folderTreeItem.SelectionItem == null && include)
        {
            // node is checked for the first time, add a selection item
            folderTreeItem.SelectionItem = new Item() { TreeItem = folderTreeItem, Include = include };
            items.Add(folderTreeItem.Path, folderTreeItem.SelectionItem);

            OnChanged();
        }
        else if (folderTreeItem.SelectionItem != null && folderTreeItem.SelectionItem.Include != include)
        {
            // node with a current state is checked/unchecked
            if (!include && (folderTreeItem.Parent == null || folderTreeItem.Parent.SelectionItem == null))
            {
                // special behavior if root node is unchecked
                items.Remove(folderTreeItem.Path);
                folderTreeItem.SelectionItem = null;
            }
            else
            {
                // determine if this node itself was specifically checked/unchecked
                var match = items.Where(x => x.Key.Equals(folderTreeItem.Path)).SingleOrDefault().Value;
                if (match != null)
                {
                    // if so, adjust it's state
                    match.Include = include;
                }
                else
                {
                    // if not, add a new selection item
                    folderTreeItem.SelectionItem = new Item() { TreeItem = folderTreeItem, Include = include };
                    items.Add(folderTreeItem.Path, folderTreeItem.SelectionItem);
                }
            }
            OnChanged();
        }

        // process children nodes if any are loaded
        if (folderTreeItem.FoldersLoaded)
        {
            foreach (var childItem in folderTreeItem.Folders)
            {
                // all children node states should be wiped out and given a reference to this nodes state
                if (items.Remove(childItem.Path))
                {
                    OnChanged();
                }
                childItem.SelectionItem = folderTreeItem.SelectionItem;
                childItem.IsChecked = include;

                // recurse to go as many levels deep as we can
                // recursive calls will skip the code above because the selection item is already set here
                Add(childItem, include);
            }
        }

        // give the parent nodes either checked or indeterminate states
        ProcessParents(folderTreeItem);
    }

    private void OnChanged()
    {
        OnPropertyChanged("SelectedFolders");
        if (Changed != null)
        {
            Changed(this, new EventArgs());
        }
    }
}

MainWindow-1.xaml

<StatusBar DockPanel.Dock="Bottom" HorizontalContentAlignment="Stretch">
    <StatusBarItem Width="Auto">
        <TextBlock x:Name="status" />
    </StatusBarItem>
    <StatusBarItem HorizontalAlignment="Right">
        <TextBlock x:Name="folderCount" Text="{Binding ElementName=folderTree, Path=Selection.FolderCount}" />
    </StatusBarItem>
</StatusBar>

MainWindow.xaml

<Window x:Class="Demo.MainWindow"
        xmlns="http://schemas.microsoft.com/winfx/2006/xaml/presentation"
        xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml"
        xmlns:ftc="clr-namespace:GeekJ.FolderTreeControl;assembly=GeekJ.FolderTreeControl"
        Title="MainWindow" Height="350" Width="525">
    <DockPanel LastChildFill="True">
        <StatusBar DockPanel.Dock="Bottom" HorizontalContentAlignment="Stretch">
            <StatusBarItem Width="Auto">
                <TextBlock x:Name="status" />
            </StatusBarItem>
        </StatusBar>
        <Grid>
            <ftc:FolderTree x:Name="folderTree" SelectionChanged="FolderTreeSelection_Changed" />
        </Grid>
    </DockPanel>
</Window>

MainWindow.xaml.cs

public partial class MainWindow : Window
{
    public BackgroundWorker worker;

    public MainWindow()
    {
        InitializeComponent();

        worker = new BackgroundWorker();
        worker.WorkerSupportsCancellation = true;

        worker.DoWork += new DoWorkEventHandler(worker_DoWork);
        worker.RunWorkerCompleted += new RunWorkerCompletedEventHandler(worker_RunWorkerCompleted);
    }

    void worker_RunWorkerCompleted(object sender, RunWorkerCompletedEventArgs e)
    {
        if (e.Cancelled)
        {
            // finished as a result of cancellation, which can only happen if the selection changed, so start over
            worker.RunWorkerAsync(folderTree.Selection);
        }
        status.Dispatcher.Invoke(new Action(delegate()
        {
            status.Text = string.Empty;
        }));
    }

    void worker_DoWork(object sender, DoWorkEventArgs e)
    {
        FolderTreeSelection selection = (FolderTreeSelection)e.Argument;
        foreach (var folder in selection.SelectedFolders)
        {
            // use the dispatcher to set the status text since we're in a different thread
            status.Dispatcher.Invoke(new Action(delegate()
            {
                status.Text = folder.FullName;
            }));

            // check the cancel flag after each iteration
            if (worker.CancellationPending)
            {
                e.Cancel = true;
                break;
            }
        }
    }

    private void FolderTreeSelection_Changed(object sender, EventArgs e)
    {
        // this will be triggered every time the selection changes
        if (worker.IsBusy)
        {
            // if the background process is already running, cancel it since the selection is different
            worker.CancelAsync();
        }
        else
        {
            // otherwise, start the background process
            worker.RunWorkerAsync(sender);
        }
    }
}