1 files changed, 88 insertions, 6 deletions

diff --git a/Modules/_testembed.c b/Modules/_testembed.c
index 51b439f..a21d251 100644
--- a/Modules/_testembed.c
+++ b/Modules/_testembed.c
@@ -1,7 +1,26 @@
 #include <Python.h>
 #include <stdio.h>
 
-void print_subinterp(void)
+/*********************************************************
+ * Embedded interpreter tests that need a custom exe
+ *
+ * Executed via 'EmbeddingTests' in Lib/test/test_capi.py
+ *********************************************************/
+
+static void _testembed_Py_Initialize(void)
+{
+    /* HACK: the "./" at front avoids a search along the PATH in
+       Modules/getpath.c */
+    Py_SetProgramName(L"./_testembed");
+    Py_Initialize();
+}
+
+
+/*****************************************************
+ * Test repeated initalisation and subinterpreters
+ *****************************************************/
+
+static void print_subinterp(void)
 {
     /* Just output some debug stuff */
     PyThreadState *ts = PyThreadState_Get();
@@ -14,7 +33,7 @@ void print_subinterp(void)
     );
 }
 
-int main(int argc, char *argv[])
+static void test_repeated_init_and_subinterpreters(void)
 {
     PyThreadState *mainstate, *substate;
 #ifdef WITH_THREAD
@@ -24,10 +43,7 @@ int main(int argc, char *argv[])
 
     for (i=0; i<3; i++) {
         printf("--- Pass %d ---\n", i);
-        /* HACK: the "./" at front avoids a search along the PATH in
-           Modules/getpath.c */
-        Py_SetProgramName(L"./_testembed");
-        Py_Initialize();
+        _testembed_Py_Initialize();
         mainstate = PyThreadState_Get();
 
 #ifdef WITH_THREAD
@@ -54,5 +70,71 @@ int main(int argc, char *argv[])
         PyEval_RestoreThread(mainstate);
         Py_Finalize();
     }
+}
+
+/*****************************************************
+ * Test forcing a particular IO encoding
+ *****************************************************/
+
+static void check_stdio_details(const char *encoding, const char * errors)
+{
+    /* Output info for the test case to check */
+    if (encoding) {
+        printf("Expected encoding: %s\n", encoding);
+    } else {
+        printf("Expected encoding: default\n");
+    }
+    if (errors) {
+        printf("Expected errors: %s\n", errors);
+    } else {
+        printf("Expected errors: default\n");
+    }
+    fflush(stdout);
+    /* Force the given IO encoding */
+    Py_SetStandardStreamEncoding(encoding, errors);
+    _testembed_Py_Initialize();
+    PyRun_SimpleString(
+        "import sys;"
+        "print('stdin: {0.encoding}:{0.errors}'.format(sys.stdin));"
+        "print('stdout: {0.encoding}:{0.errors}'.format(sys.stdout));"
+        "print('stderr: {0.encoding}:{0.errors}'.format(sys.stderr));"
+        "sys.stdout.flush()"
+    );
+    Py_Finalize();
+}
+
+static void test_forced_io_encoding(void)
+{
+    /* Check various combinations */
+    printf("--- Use defaults ---\n");
+    check_stdio_details(NULL, NULL);
+    printf("--- Set errors only ---\n");
+    check_stdio_details(NULL, "surrogateescape");
+    printf("--- Set encoding only ---\n");
+    check_stdio_details("latin-1", NULL);
+    printf("--- Set encoding and errors ---\n");
+    check_stdio_details("latin-1", "surrogateescape");
+
+    /* Check calling after initialization fails */
+    Py_Initialize();
+
+    if (Py_SetStandardStreamEncoding(NULL, NULL) == 0) {
+        printf("Unexpected success calling Py_SetStandardStreamEncoding");
+    }
+    Py_Finalize();
+}
+
+/* Different embedding tests */
+int main(int argc, char *argv[])
+{
+
+    /* TODO: Check the argument string to allow for more test cases */
+    if (argc > 1) {
+        /* For now: assume "forced_io_encoding */
+        test_forced_io_encoding();
+    } else {
+        /* Run the original embedding test case by default */
+        test_repeated_init_and_subinterpreters();
+    }
     return 0;
 }