From 8e3f0715028cbaf9360cd10f395be07fbb651dec Mon Sep 17 00:00:00 2001
From: John Haddon <thehaddonyoof@gmail.com>
Date: Thu, 14 Sep 2023 12:21:14 +0100
Subject: [PATCH] USDScene : Support reading from in-memory stages via
 UsdUtilsStageCache

---
 Changes                                       |  1 +
 SConstruct                                    |  1 +
 contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp  | 20 +++++++++++++++++--
 .../IECoreUSD/test/IECoreUSD/USDSceneTest.py  | 10 ++++++++++
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/Changes b/Changes
index 932a206b63..3f1175ced4 100644
--- a/Changes
+++ b/Changes
@@ -5,6 +5,7 @@ Improvements
 ------------
 
 - USD PrimitiveAlgo : Added `readPrimitiveVariable()` utility method for reading from regular `UsdAttributes`.
+- USDScene : Added support for reading from in-memory stages by passing a filename of the form `stageCache:{id}.usd` where `{id}` specifies a stage which has been inserted in the `UsdUtilsStageCache`.
 
 Fixes
 -----
diff --git a/SConstruct b/SConstruct
index d410a5268c..71a02ef57c 100644
--- a/SConstruct
+++ b/SConstruct
@@ -2967,6 +2967,7 @@ else :
 		"usdLux",
 		"usdSkel",
 		"usdShade",
+		"usdUtils",
 		"sdf",
 		"tf",
 		"pcp",
diff --git a/contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp b/contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp
index 9c17fd9c42..f9aa69e8fd 100644
--- a/contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp
+++ b/contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp
@@ -70,6 +70,7 @@ IECORE_PUSH_DEFAULT_VISIBILITY
 #include "pxr/usd/usdShade/material.h"
 #include "pxr/usd/usdShade/materialBindingAPI.h"
 #include "pxr/usd/usdShade/connectableAPI.h"
+#include "pxr/usd/usdUtils/stageCache.h"
 #ifdef IECOREUSD_WITH_OPENVDB
 #include "pxr/usd/usdVol/fieldBase.h"
 #endif
@@ -84,6 +85,7 @@ IECORE_POP_DEFAULT_VISIBILITY
 
 #include "tbb/concurrent_hash_map.h"
 
+#include <filesystem>
 #include <iostream>
 #include <mutex>
 
@@ -688,10 +690,24 @@ class USDScene::IO : public RefCounted
 			switch( openMode )
 			{
 				case IndexedIO::Read : {
-					pxr::UsdStageRefPtr stage = pxr::UsdStage::Open( fileName );
+					static const std::string g_stageCachePrefix( "stageCache:" );
+					pxr::UsdStageRefPtr stage;
+					if( boost::starts_with( fileName, g_stageCachePrefix ) )
+					{
+						// Get Id from filename of form "stageCache:{id}.usd"
+						std::filesystem::path path( fileName.substr( g_stageCachePrefix.size() ) );
+						path.replace_extension();
+						stage = pxr::UsdUtilsStageCache::Get().Find(
+							pxr::UsdStageCache::Id::FromString( path.string() )
+						);
+					}
+					else
+					{
+						stage = pxr::UsdStage::Open( fileName );
+					}
 					if( !stage )
 					{
-						throw IECore::Exception( boost::str( boost::format( "USDScene : Failed to open USD file: '%1%'" ) % fileName ) );
+						throw IECore::Exception( boost::str( boost::format( "USDScene : Failed to open USD stage : '%1%'" ) % fileName ) );
 					}
 					return stage;
 				}
diff --git a/contrib/IECoreUSD/test/IECoreUSD/USDSceneTest.py b/contrib/IECoreUSD/test/IECoreUSD/USDSceneTest.py
index 797a262b43..e52717d613 100644
--- a/contrib/IECoreUSD/test/IECoreUSD/USDSceneTest.py
+++ b/contrib/IECoreUSD/test/IECoreUSD/USDSceneTest.py
@@ -3911,5 +3911,15 @@ def testUnconnectedMaterialOutput( self ) :
 		self.assertNotIn( "cycles:surface", sphere.attributeNames() )
 		self.assertIsNone( sphere.readAttribute( "cycles:surface", 0 ) )
 
+	def testReadFromStageCache( self ) :
+
+		stage = pxr.Usd.Stage.CreateInMemory()
+		pxr.UsdGeom.Sphere.Define( stage, "/sphere" )
+		id = pxr.UsdUtils.StageCache.Get().Insert( stage )
+
+		root = IECoreScene.SceneInterface.create( "stageCache:{}.usd".format( id.ToString() ), IECore.IndexedIO.OpenMode.Read )
+		self.assertEqual( root.childNames(), [ "sphere" ] )
+		self.assertIsInstance( root.child( "sphere" ).readObject( 0 ), IECoreScene.SpherePrimitive )
+
 if __name__ == "__main__":
 	unittest.main()