Fixed cached generators being exhausted after the first iteration.

Author: jordanbriere

addons/source-python/packages/source-python/core/cache.py

@@ -12,6 +12,7 @@
 # =============================================================================
 # Source.Python Imports
 #   Core
+from _core._cache import CachedGenerator
 from _core._cache import CachedProperty
 from _core._cache import cached_property
 
@@ -20,6 +21,7 @@
 # >> ALL DECLARATION
 # =============================================================================
 __all__ = [
+    'CachedGenerator',
     'CachedProperty',
     'cached_property'
 ]

src/core/modules/core/core_cache.cpp

@@ -60,6 +60,13 @@ object CCachedProperty::_callable_check(object function, const char *szName)
 	return function;
 }
 
+object CCachedProperty::_prepare_value(object value)
+{
+	if (PyGen_Check(value.ptr()))
+		value = object(CCachedGenerator(value));
+	return value;
+}
+
 
 object CCachedProperty::get_getter()
 {
@@ -146,9 +153,11 @@ object CCachedProperty::__get__(object instance, object owner=object())
 		);
 
 	dict cache = extract<dict>(pCache);
-	cache[m_name] = m_fget(
-		*(make_tuple(handle<>(borrowed(instance.ptr()))) + m_args),
-		**m_kwargs
+	cache[m_name] = _prepare_value(
+		m_fget(
+			*(make_tuple(handle<>(borrowed(instance.ptr()))) + m_args),
+			**m_kwargs
+		)
 	);
 
 	return cache[m_name];
@@ -170,7 +179,7 @@ void CCachedProperty::__set__(object instance, object value)
 
 	dict cache = extract<dict>(instance.attr("__dict__"));
 	if (!result.is_none())
-		cache[m_name] = result;
+		cache[m_name] = _prepare_value(result);
 	else
 		PyDict_DelItemString(cache.ptr(), extract<const char *>(m_name));
 }
@@ -209,3 +218,57 @@ CCachedProperty *CCachedProperty::wrap_descriptor(object descriptor, object owne
 
 	return pProperty;
 }
+
+
+//-----------------------------------------------------------------------------
+// CCachedGenerator class.
+//-----------------------------------------------------------------------------
+CCachedGenerator::CCachedGenerator(object generator)
+{
+	if (!PyGen_Check(generator.ptr()))
+		BOOST_RAISE_EXCEPTION(
+			PyExc_TypeError,
+			"The given generator is invalid."
+		);
+
+	object frame = generator.attr("gi_frame");
+	if (frame.is_none())
+		BOOST_RAISE_EXCEPTION(
+			PyExc_ValueError,
+			"The given generator is exhausted."
+		);
+
+	m_generator = generator;
+}
+
+
+object CCachedGenerator::get_generator()
+{
+	return m_generator;
+}
+
+
+object CCachedGenerator::__iter__()
+{
+	return m_generator.is_none() ? m_generated_values.attr("__iter__")() : object(ptr(this));
+}
+
+
+object CCachedGenerator::__next__()
+{
+	object value;
+	if (!m_generator.is_none())
+	{
+		try
+		{
+			value = m_generator.attr("__next__")();
+			m_generated_values.append(value);
+		}
+		catch(...)
+		{
+			m_generator = object();
+			BOOST_RAISE_EXCEPTION(PyExc_StopIteration, "StopIteration");
+		}
+	}
+	return value;
+}

src/core/modules/core/core_cache.h

@@ -43,6 +43,7 @@ class CCachedProperty
 	CCachedProperty(object fget, object fset, object fdel, const char *doc, boost::python::tuple args, dict kwargs);
 
 	static object _callable_check(object function, const char *szName);
+	static object _prepare_value(object value);
 
 	object get_getter();
 	object set_getter(object fget);
@@ -81,4 +82,23 @@ class CCachedProperty
 };
 
 
+//-----------------------------------------------------------------------------
+// CCachedGenerator class.
+//-----------------------------------------------------------------------------
+class CCachedGenerator
+{
+public:
+	CCachedGenerator(object generator);
+
+	object get_generator();
+
+	object __iter__();
+	object __next__();
+
+private:
+	object m_generator;
+	list m_generated_values;
+};
+
+
 #endif // _CORE_CACHE_H

src/core/modules/core/core_cache_wrap.cpp

@@ -36,6 +36,7 @@
 // Forward declarations.
 //-----------------------------------------------------------------------------
 static void export_cached_property(scope);
+static void export_cached_generator(scope);
 
 
 //-----------------------------------------------------------------------------
@@ -44,6 +45,7 @@ static void export_cached_property(scope);
 DECLARE_SP_SUBMODULE(_core, _cache)
 {
 	export_cached_property(_cache);
+	export_cached_generator(_cache);
 }
 
 
@@ -279,3 +281,43 @@ void export_cached_property(scope _cache)
 
 	scope().attr("cached_property") = scope().attr("CachedProperty");
 }
+
+
+//-----------------------------------------------------------------------------
+// Exports CCachedGenerator.
+//-----------------------------------------------------------------------------
+void export_cached_generator(scope _cache)
+{
+	class_<CCachedGenerator, CCachedGenerator *> CachedGenerator("CachedGenerator",
+		init<object>(
+			(
+				arg("generator")
+			),
+			"Represents a cached generator.\n"
+			"If a :class:`core.cache.CachedProperty` returns a generator, it"
+			" is being cached as an instance of this class. Then, when"
+			" this instance is iterated over for the first time, the original"
+			" generator is processed and the generated values are cached.\n"
+			"\n"
+			":param generator generator:\n"
+			"	The wrapped generator instance.\n"
+			"\n"
+			":raises TypeError:\n"
+			"	If the given generator is invalid.\n"
+			":raises ValueError:\n"
+			"	If the given generator is exhausted."
+		)
+	);
+
+	CachedGenerator.def(
+		"__iter__",
+		&CCachedGenerator::__iter__,
+		"Returns an iterator iterating over the generated values of the wrapped generator."
+	);
+
+	CachedGenerator.def(
+		"__next__",
+		&CCachedGenerator::__next__,
+		"Returns the next value from the current iteration."
+	);
+}