1 files changed, 25 insertions, 16 deletions

diff --git a/unix/main.c b/unix/main.c
index d916cdeb88..01ebdce175 100644
--- a/unix/main.c
+++ b/unix/main.c
@@ -90,19 +90,33 @@ STATIC int handle_uncaught_exception(mp_obj_base_t *exc) {
     return 1;
 }
 
+#define LEX_SRC_STR (1)
+#define LEX_SRC_VSTR (2)
+#define LEX_SRC_FILENAME (3)
+#define LEX_SRC_STDIN (4)
+
 // Returns standard error codes: 0 for success, 1 for all other errors,
 // except if FORCED_EXIT bit is set then script raised SystemExit and the
 // value of the exit is in the lower 8 bits of the return value
-STATIC int execute_from_lexer(mp_lexer_t *lex, mp_parse_input_kind_t input_kind, bool is_repl) {
-    if (lex == NULL) {
-        printf("MemoryError: lexer could not allocate memory\n");
-        return 1;
-    }
-
+STATIC int execute_from_lexer(int source_kind, const void *source, mp_parse_input_kind_t input_kind, bool is_repl) {
     mp_hal_set_interrupt_char(CHAR_CTRL_C);
 
     nlr_buf_t nlr;
     if (nlr_push(&nlr) == 0) {
+        // create lexer based on source kind
+        mp_lexer_t *lex;
+        if (source_kind == LEX_SRC_STR) {
+            const char *line = source;
+            lex = mp_lexer_new_from_str_len(MP_QSTR__lt_stdin_gt_, line, strlen(line), false);
+        } else if (source_kind == LEX_SRC_VSTR) {
+            const vstr_t *vstr = source;
+            lex = mp_lexer_new_from_str_len(MP_QSTR__lt_stdin_gt_, vstr->buf, vstr->len, false);
+        } else if (source_kind == LEX_SRC_FILENAME) {
+            lex = mp_lexer_new_from_file((const char*)source);
+        } else { // LEX_SRC_STDIN
+            lex = mp_lexer_new_from_fd(MP_QSTR__lt_stdin_gt_, 0, false);
+        }
+
         qstr source_name = lex->source_name;
 
         #if MICROPY_PY___FILE__
@@ -240,8 +254,7 @@ STATIC int do_repl(void) {
 
         mp_hal_stdio_mode_orig();
 
-        mp_lexer_t *lex = mp_lexer_new_from_str_len(MP_QSTR__lt_stdin_gt_, line.buf, line.len, false);
-        ret = execute_from_lexer(lex, parse_input_kind, true);
+        ret = execute_from_lexer(LEX_SRC_VSTR, &line, parse_input_kind, true);
         if (ret & FORCED_EXIT) {
             return ret;
         }
@@ -268,8 +281,7 @@ STATIC int do_repl(void) {
             line = line3;
         }
 
-        mp_lexer_t *lex = mp_lexer_new_from_str_len(MP_QSTR__lt_stdin_gt_, line, strlen(line), false);
-        int ret = execute_from_lexer(lex, MP_PARSE_SINGLE_INPUT, true);
+        int ret = execute_from_lexer(LEX_SRC_STR, line, MP_PARSE_SINGLE_INPUT, true);
         if (ret & FORCED_EXIT) {
             return ret;
         }
@@ -280,13 +292,11 @@ STATIC int do_repl(void) {
 }
 
 STATIC int do_file(const char *file) {
-    mp_lexer_t *lex = mp_lexer_new_from_file(file);
-    return execute_from_lexer(lex, MP_PARSE_FILE_INPUT, false);
+    return execute_from_lexer(LEX_SRC_FILENAME, file, MP_PARSE_FILE_INPUT, false);
 }
 
 STATIC int do_str(const char *str) {
-    mp_lexer_t *lex = mp_lexer_new_from_str_len(MP_QSTR__lt_stdin_gt_, str, strlen(str), false);
-    return execute_from_lexer(lex, MP_PARSE_FILE_INPUT, false);
+    return execute_from_lexer(LEX_SRC_STR, str, MP_PARSE_FILE_INPUT, false);
 }
 
 STATIC int usage(char **argv) {
@@ -585,8 +595,7 @@ MP_NOINLINE int main_(int argc, char **argv) {
             ret = do_repl();
             prompt_write_history();
         } else {
-            mp_lexer_t *lex = mp_lexer_new_from_fd(MP_QSTR__lt_stdin_gt_, 0, false);
-            ret = execute_from_lexer(lex, MP_PARSE_FILE_INPUT, false);
+            ret = execute_from_lexer(LEX_SRC_STDIN, NULL, MP_PARSE_FILE_INPUT, false);
         }
     }