From 5feecd2243811fb899d056ecdcd2e962765e501a Mon Sep 17 00:00:00 2001
From: rdb <git@rdb.name>
Date: Thu, 28 Apr 2016 11:08:23 +0200
Subject: [PATCH] Fix iterator support in interrogate

---
 .../interfaceMakerPythonNative.cxx            | 21 +++++++++++++++----
 .../interrogate/interfaceMakerPythonNative.h  |  5 ++++-
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/dtool/src/interrogate/interfaceMakerPythonNative.cxx b/dtool/src/interrogate/interfaceMakerPythonNative.cxx
index eba00431d3..5dde93b7d1 100644
--- a/dtool/src/interrogate/interfaceMakerPythonNative.cxx
+++ b/dtool/src/interrogate/interfaceMakerPythonNative.cxx
@@ -1738,7 +1738,7 @@ write_module_class(ostream &out, Object *obj) {
 
       switch (rfi->second._wrapper_type) {
       case WT_no_params:
-      case WT_iter_next: // TODO: fix iter_next to return NULL instead of None
+      case WT_iter_next:
         // PyObject *func(PyObject *self)
         {
           out << "//////////////////\n";
@@ -1751,9 +1751,15 @@ write_module_class(ostream &out, Object *obj) {
           out << "    return NULL;\n";
           out << "  }\n\n";
 
+          int return_flags = RF_pyobject | RF_err_null;
+          if (rfi->second._wrapper_type == WT_iter_next) {
+            // If the function returns NULL, we should return NULL to indicate
+            // a StopIteration, rather than returning None.
+            return_flags |= RF_preserve_null;
+          }
           string expected_params;
           write_function_forset(out, def._remaps, 0, 0, expected_params, 2, true, true,
-                                AT_no_args, RF_pyobject | RF_err_null, false);
+                                AT_no_args, return_flags, false);
 
           out << "  if (!_PyErr_OCCURRED()) {\n";
           out << "    return Dtool_Raise_BadArgumentsError(\n";
@@ -5856,8 +5862,15 @@ write_function_instance(ostream &out, FunctionRemap *remap,
       indent(out, indent_level) << "Py_INCREF(Py_None);\n";
       indent(out, indent_level) << "return Py_None;\n";
 
+    } else if (return_flags & RF_preserve_null) {
+      indent(out, indent_level) << "if (" << return_expr << " == NULL) {\n";
+      indent(out, indent_level) << "  return NULL;\n";
+      indent(out, indent_level) << "} else {\n";
+      pack_return_value(out, indent_level + 2, remap, return_expr, return_flags);
+      indent(out, indent_level) << "}\n";
+
     } else {
-      pack_return_value(out, indent_level, remap, return_expr);
+      pack_return_value(out, indent_level, remap, return_expr, return_flags);
     }
 
   } else if (return_flags & RF_coerced) {
@@ -6014,7 +6027,7 @@ error_raise_return(ostream &out, int indent_level, int return_flags,
  */
 void InterfaceMakerPythonNative::
 pack_return_value(ostream &out, int indent_level, FunctionRemap *remap,
-                  string return_expr) {
+                  string return_expr, int return_flags) {
 
   ParameterRemap *return_type = remap->_return_type;
   CPPType *orig_type = return_type->get_orig_type();
diff --git a/dtool/src/interrogate/interfaceMakerPythonNative.h b/dtool/src/interrogate/interfaceMakerPythonNative.h
index 1c7b7d98de..a571e3323f 100644
--- a/dtool/src/interrogate/interfaceMakerPythonNative.h
+++ b/dtool/src/interrogate/interfaceMakerPythonNative.h
@@ -101,6 +101,9 @@ private:
     // Assign to the coerced argument, in the case of a coercion constructor.
     RF_coerced = 0x040,
 
+    // Don't automatically map NULL to None
+    RF_preserve_null = 0x080,
+
     // These indicate what should be returned on error.
     RF_err_notimplemented = 0x002,
     RF_err_null = 0x004,
@@ -164,7 +167,7 @@ private:
                           const string &exc_type, const string &message,
                           const string &format_args = "");
   void pack_return_value(ostream &out, int indent_level, FunctionRemap *remap,
-                         std::string return_expr);
+                         std::string return_expr, int return_flags);
 
   void write_make_seq(ostream &out, Object *obj, const std::string &ClassName,
                       const std::string &cClassName, MakeSeq *make_seq);