Generic SceneHierarchyWindow static utility methods.

SceneHierarchyWindowUtility.cs

using System.Collections.Generic;
using System.Reflection;

using UnityEngine;

using UnityEditor;
using UnityEditor.IMGUI.Controls;

public static class SceneHierarchyWindowUtility
{
	private static readonly System.Type TypeOfSceneHierarchyWindow = null;
	private static readonly FieldInfo SceneHierarchyWindowTreeViewControllerField = null;
	private static readonly MethodInfo SceneHierarchyWindowSetExpandedRecursiveMethod = null;

	private static readonly MethodInfo TreeViewControllerFindItemMethod = null;
	private static readonly MethodInfo TreeViewControllerGetSelectionMethod = null;

	private static readonly System.Type TypeOfGameObjectTreeViewItem = null;
	private static readonly PropertyInfo GameObjectTreeViewItemIsSceneHeader = null;

	static SceneHierarchyWindowUtility()
	{
		// Fetch SceneHierarchyWindow type and related members through reflection.
		// We get the correctly loaded assembly (UnityEditor) through a visible type (Editor) because its easier to read and more compact that way.
		TypeOfSceneHierarchyWindow = typeof(Editor).Assembly.GetType("UnityEditor.SceneHierarchyWindow");
		SceneHierarchyWindowTreeViewControllerField = TypeOfSceneHierarchyWindow.GetField("m_TreeView", BindingFlags.NonPublic | BindingFlags.Instance);
		SceneHierarchyWindowSetExpandedRecursiveMethod = TypeOfSceneHierarchyWindow.GetMethod("SetExpandedRecursive", BindingFlags.Public | BindingFlags.Instance);

		// Fetch GameObjectTreeViewItem type and related members through reflection.
		TypeOfGameObjectTreeViewItem = typeof(Editor).Assembly.GetType("UnityEditor.GameObjectTreeViewItem");
		GameObjectTreeViewItemIsSceneHeader = TypeOfGameObjectTreeViewItem.GetProperty("isSceneHeader", BindingFlags.Public | BindingFlags.Instance);

		// Fetch TreeViewController type and related members through reflection.
		System.Type TypeOfTreeViewController = typeof(TreeView).Assembly.GetType("UnityEditor.IMGUI.Controls.TreeViewController");
		TreeViewControllerFindItemMethod = TypeOfTreeViewController.GetMethod("FindItem", BindingFlags.Public | BindingFlags.Instance);
		TreeViewControllerGetSelectionMethod = TypeOfTreeViewController.GetMethod("GetSelection", BindingFlags.Public | BindingFlags.Instance);
	}

	[MenuItem("Tools/Collapse Scenes With Selection")]
	public static void CollapseScenesWithSelection()
	{
		// Fetch every SceneHierarchyWindow currently loaded.
		Object[] sceneHierarchyWindows = Resources.FindObjectsOfTypeAll(TypeOfSceneHierarchyWindow);

		foreach(Object sceneHierarchyWindow in sceneHierarchyWindows)
		{
			// Get the id(s) of the current TreeViewItem(s) selected in this SceneHierarchyWindow.
			int[] selectedIds = GetSelection(sceneHierarchyWindow);
			List<int> selectedScenesIds = new List<int>();

			foreach(int selectedId in selectedIds)
			{
				// Get the TreeViewItem that represent the selected object inside the current hierarchy (GameObject, Scene or root).
				TreeViewItem treeViewItem = FindItem(sceneHierarchyWindow, selectedId);

				// Get the TreeViewItem that represent the root Scene of this item (it can be itself).
				treeViewItem = GetSceneItem(treeViewItem);

				// This is the tricky part: the id store inside the TreeViewItem represent the handle for this scene.
				// This correlation is just a practical trick (found while looking inside Unity's Editor code).
				// That way we don't have to fetch the Scene object (because we already have what we want).
				if(!selectedScenesIds.Contains(treeViewItem.id))
				{
					selectedScenesIds.Add(treeViewItem.id);
				}
			}

			foreach(int selectedSceneId in selectedScenesIds)
			{
				CollapseSceneHierarchy(sceneHierarchyWindow, selectedSceneId);
			}

			EditorWindow editorWindow = (EditorWindow)sceneHierarchyWindow;
			editorWindow.Repaint();
		}
	}

	private static void CollapseSceneHierarchy(object sceneHierarchyWindow, int sceneId)
	{
		object[] parameters = new object[2]
		{
			sceneId,
			false
		};

		SceneHierarchyWindowSetExpandedRecursiveMethod.Invoke(sceneHierarchyWindow, parameters);
	}

	private static TreeViewItem GetSceneItem(TreeViewItem treeViewItem)
	{
		while(treeViewItem != null)
		{
			if(TreeViewItemIsSceneHeader(treeViewItem))
			{
				return treeViewItem;
			}

			treeViewItem = treeViewItem.parent;
		}

		return null;
	}

	private static bool TreeViewItemIsSceneHeader(TreeViewItem treeViewItem)
	{
		System.Type treeViewItemType = treeViewItem.GetType();
		if(TypeOfGameObjectTreeViewItem.IsAssignableFrom(treeViewItemType))
		{
			return (bool)GameObjectTreeViewItemIsSceneHeader.GetValue(treeViewItem, null);
		}

		return false;
	}

	private static object GetTreeViewController(object sceneHierarchyWindow)
	{
		return SceneHierarchyWindowTreeViewControllerField.GetValue(sceneHierarchyWindow);
	}

	private static int[] GetSelection(object sceneHierarchyWindow)
	{
		object treeViewController = GetTreeViewController(sceneHierarchyWindow);
		return (int[])TreeViewControllerGetSelectionMethod.Invoke(treeViewController, null);
	}

	private static TreeViewItem FindItem(object sceneHierarchyWindow, int id)
	{
		object treeViewController = GetTreeViewController(sceneHierarchyWindow);
		object[] parameters = new object[1]
		{
			id
		};

		return TreeViewControllerFindItemMethod.Invoke(treeViewController, parameters) as TreeViewItem;
	}
}