From be721f34157672bc4a5eebe9e908458e250ef1ce Mon Sep 17 00:00:00 2001
From: Krasimir Angelov <kr.angelov@gmail.com>
Date: Sat, 8 Oct 2022 17:21:36 +0200
Subject: [PATCH] PGF.embed can now augment existing modules

---
 src/runtime/python/pypgf.c | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/runtime/python/pypgf.c b/src/runtime/python/pypgf.c
index fb592792c..64ab9946c 100644
--- a/src/runtime/python/pypgf.c
+++ b/src/runtime/python/pypgf.c
@@ -399,9 +399,26 @@ pgf_embed_funs(PgfItor* fn, PgfText* name, object value, PgfExn* err)
 static PyObject*
 PGF_embed(PGFObject* self, PyObject *modname)
 {
-	PyObject *m = PyImport_AddModuleObject(modname);
-	if (m == NULL)
-		return NULL;
+    PyObject *m = PyImport_Import(modname);
+    if (m == NULL) {
+        PyObject *globals = PyEval_GetGlobals();
+        if (globals != NULL) {
+            PyObject *exc = PyDict_GetItemString(globals, "ModuleNotFoundError");
+            if (exc != NULL) {
+                if (PyErr_ExceptionMatches(exc)) {
+                    PyErr_Clear();
+                    m = PyImport_AddModuleObject(modname);
+                    Py_INCREF(m);
+                }
+            }
+        }
+    }
+    if (m == NULL)
+        return NULL;
+
+    if (PyModule_AddObjectRef(m, "__pgf__", (PyObject*) self) != 0) {
+        return NULL;
+    }
 
     PgfExn err;
     PyPGFClosure clo = { { pgf_embed_funs }, self, m };
@@ -411,7 +428,6 @@ PGF_embed(PGFObject* self, PyObject *modname)
         return NULL;
     }
 
-    Py_INCREF(m);
     return m;
 }