Fix call_parse to reliably detect direct script execution

fastcore/_modidx.py

@@ -625,7 +625,9 @@
                                  'fastcore.script._HelpFormatter.__init__': ('script.html#_helpformatter.__init__', 'fastcore/script.py'),
                                  'fastcore.script._HelpFormatter._expand_help': ( 'script.html#_helpformatter._expand_help',
                                                                                   'fastcore/script.py'),
+                                 'fastcore.script._is_script_run': ('script.html#_is_script_run', 'fastcore/script.py'),
                                  'fastcore.script._is_union': ('script.html#_is_union', 'fastcore/script.py'),
+                                 'fastcore.script._run_cli': ('script.html#_run_cli', 'fastcore/script.py'),
                                  'fastcore.script._union_parser': ('script.html#_union_parser', 'fastcore/script.py'),
                                  'fastcore.script._union_type': ('script.html#_union_type', 'fastcore/script.py'),
                                  'fastcore.script.anno_parser': ('script.html#anno_parser', 'fastcore/script.py'),

fastcore/script.py

@@ -149,32 +149,42 @@ def set_ctx(cv, val=True):
     finally: cv.reset(token)
 
 # %% ../nbs/06_script.ipynb #fc816498
+def _is_script_run(frame):
+    "True if `frame` is the body of a file being run directly (`python foo.py`, `python -m foo`, or `%run foo.py`)"
+    g = frame.f_globals if frame else {}
+    if g.get('__name__')!='__main__' or not g.get('__file__'): return False
+    return Path(g['__file__']).resolve()==Path(frame.f_code.co_filename).resolve()
+
+# %% ../nbs/06_script.ipynb #dee5e259
 _in_call_parse = ContextVar('_in_call_parse', default=False)
 
+def _run_cli(func, nested):
+    "Parse `sys.argv` with `anno_parser` and call `func` with the result"
+    if len(sys.argv)>1 and sys.argv[1]=='': sys.argv.pop(1)
+    p = anno_parser(func)
+    if nested: args, sys.argv[1:] = p.parse_known_args()
+    else: args = p.parse_args()
+    args = args.__dict__
+    xtra = otherwise(args.pop('xtra', ''), eq(1), p.prog)
+    tfunc = trace(func) if args.pop('pdb', False) else func
+    res = tfunc(**merge(args, args_from_prog(func, xtra)))
+    return asyncio.run(res) if inspect.isawaitable(res) else res
+
 def call_parse(func=None, nested=False):
     "Decorator to create a simple CLI from `func` using `anno_parser`"
     if func is None: return partial(call_parse, nested=nested)
-
     @wraps(func)
     def _f(*args, **kwargs):
         if args or kwargs or _in_call_parse.get(): return func(*args, **kwargs)
         with set_ctx(_in_call_parse):
             mod = inspect.getmodule(inspect.currentframe().f_back)
             if not mod: return func(*args, **kwargs)
             if not SCRIPT_INFO.func and mod.__name__=="__main__": SCRIPT_INFO.func = func.__name__
-            if len(sys.argv)>1 and sys.argv[1]=='': sys.argv.pop(1)
-            p = anno_parser(func)
-            if nested: args, sys.argv[1:] = p.parse_known_args()
-            else: args = p.parse_args()
-            args = args.__dict__
-            xtra = otherwise(args.pop('xtra', ''), eq(1), p.prog)
-            tfunc = trace(func) if args.pop('pdb', False) else func
-            res = tfunc(**merge(args, args_from_prog(func, xtra)))
-            return asyncio.run(res) if inspect.isawaitable(res) else res
-
-    mod = inspect.getmodule(inspect.currentframe().f_back)
-    if getattr(mod, '__name__', '') =="__main__":
-        setattr(mod, func.__name__, _f)
+            return _run_cli(func, nested)
+
+    frame = inspect.currentframe().f_back
+    if _is_script_run(frame):
+        frame.f_globals[func.__name__] = _f
         SCRIPT_INFO.func = func.__name__
-        return _f()
+        with set_ctx(_in_call_parse): return _run_cli(func, nested)
     else: return _f

nbs/06_script.ipynb

@@ -195,7 +195,9 @@
    "outputs": [],
    "source": [
     "#| hide\n",
-    "from fastcore.test import *"
+    "from fastcore.test import *\n",
+    "\n",
+    "import runpy, tempfile"
    ]
   },
   {
@@ -371,8 +373,18 @@
      "text": [
       "Help on function f in module __main__:\n",
       "\n",
-      "f(required: int <Required param>, a: bool_arg <param 1>, b: str <param 2> = 'test')\n",
-      "    my docs\n",
+      "f(\n",
+      "    required: int <Required param>,\n",
+      "    a: bool_arg <param 1>,\n",
+      "    b: str <param 2> = 'test'\n",
+      ")\n",
+      "    my docs\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
       "\n"
      ]
     }
@@ -782,36 +794,70 @@
    "id": "fc816498",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "#| export\n",
+    "def _is_script_run(frame):\n",
+    "    \"True if `frame` is the body of a file being run directly (`python foo.py`, `python -m foo`, or `%run foo.py`)\"\n",
+    "    g = frame.f_globals if frame else {}\n",
+    "    if g.get('__name__')!='__main__' or not g.get('__file__'): return False\n",
+    "    return Path(g['__file__']).resolve()==Path(frame.f_code.co_filename).resolve()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d307a7c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _chk(co_file, name='__main__', file='ascript.py'):\n",
+    "    g = {'__name__':name, '__file__':file, 'inspect':inspect, '_is_script_run':_is_script_run}\n",
+    "    return eval(compile('_is_script_run(inspect.currentframe())', co_file, 'eval'), g)\n",
+    "\n",
+    "test_eq(_chk('ascript.py'), True)               # file run directly\n",
+    "test_eq(_chk('ascript.py', name='amod'), False) # imported module\n",
+    "test_eq(_chk('other.py'), False)                # notebook cell: code isn't __main__'s file"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dee5e259",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "#| export\n",
     "_in_call_parse = ContextVar('_in_call_parse', default=False)\n",
     "\n",
+    "def _run_cli(func, nested):\n",
+    "    \"Parse `sys.argv` with `anno_parser` and call `func` with the result\"\n",
+    "    if len(sys.argv)>1 and sys.argv[1]=='': sys.argv.pop(1)\n",
+    "    p = anno_parser(func)\n",
+    "    if nested: args, sys.argv[1:] = p.parse_known_args()\n",
+    "    else: args = p.parse_args()\n",
+    "    args = args.__dict__\n",
+    "    xtra = otherwise(args.pop('xtra', ''), eq(1), p.prog)\n",
+    "    tfunc = trace(func) if args.pop('pdb', False) else func\n",
+    "    res = tfunc(**merge(args, args_from_prog(func, xtra)))\n",
+    "    return asyncio.run(res) if inspect.isawaitable(res) else res\n",
+    "\n",
     "def call_parse(func=None, nested=False):\n",
     "    \"Decorator to create a simple CLI from `func` using `anno_parser`\"\n",
     "    if func is None: return partial(call_parse, nested=nested)\n",
-    "\n",
     "    @wraps(func)\n",
     "    def _f(*args, **kwargs):\n",
     "        if args or kwargs or _in_call_parse.get(): return func(*args, **kwargs)\n",
     "        with set_ctx(_in_call_parse):\n",
     "            mod = inspect.getmodule(inspect.currentframe().f_back)\n",
     "            if not mod: return func(*args, **kwargs)\n",
     "            if not SCRIPT_INFO.func and mod.__name__==\"__main__\": SCRIPT_INFO.func = func.__name__\n",
-    "            if len(sys.argv)>1 and sys.argv[1]=='': sys.argv.pop(1)\n",
-    "            p = anno_parser(func)\n",
-    "            if nested: args, sys.argv[1:] = p.parse_known_args()\n",
-    "            else: args = p.parse_args()\n",
-    "            args = args.__dict__\n",
-    "            xtra = otherwise(args.pop('xtra', ''), eq(1), p.prog)\n",
-    "            tfunc = trace(func) if args.pop('pdb', False) else func\n",
-    "            res = tfunc(**merge(args, args_from_prog(func, xtra)))\n",
-    "            return asyncio.run(res) if inspect.isawaitable(res) else res\n",
-    "\n",
-    "    mod = inspect.getmodule(inspect.currentframe().f_back)\n",
-    "    if getattr(mod, '__name__', '') ==\"__main__\":\n",
-    "        setattr(mod, func.__name__, _f)\n",
+    "            return _run_cli(func, nested)\n",
+    "\n",
+    "    frame = inspect.currentframe().f_back\n",
+    "    if _is_script_run(frame):\n",
+    "        frame.f_globals[func.__name__] = _f\n",
     "        SCRIPT_INFO.func = func.__name__\n",
-    "        return _f()\n",
+    "        with set_ctx(_in_call_parse): return _run_cli(func, nested)\n",
     "    else: return _f"
    ]
   },
@@ -851,15 +897,30 @@
   },
   {
    "cell_type": "markdown",
-   "id": "8ec22c23",
+   "id": "2e207b67",
    "metadata": {},
    "source": [
-    "This is the main way to use `fastcore.script`; decorate your function with `call_parse`, add `Param` annotations (as shown above) or type annotations and docments, and it can then be used as a script."
+    "Calling a decorated function from Python (a notebook, the REPL, or another decorated function) is a plain call, and `sys.argv` is ignored; argv is only parsed when the function runs as a CLI entry point:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6b62dc1b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@call_parse\n",
+    "def main(a:int=0): return test_add(a, a)\n",
+    "\n",
+    "argv,sys.argv = sys.argv,['prog', '--a', '3']\n",
+    "try: test_eq(main(), 0)  # interactive call: uses defaults, argv untouched\n",
+    "finally: sys.argv = argv"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "9c45d452",
+   "id": "cbb6d241",
    "metadata": {},
    "source": [
     "Use the `nested` keyword argument to create nested parsers, where earlier parsers consume only their known args from `sys.argv` before later parsers are used. This is useful to create one command line application that executes another. For example:\n",
@@ -877,6 +938,56 @@
     "would display `myrunner`'s help and not `script.py`'s."
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f22c7158",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "argv,sys.argv = sys.argv,['runner', '--msg', 'hello']\n",
+    "with tempfile.TemporaryDirectory() as d:\n",
+    "    p,out = Path(d)/'inner.py', Path(d)/'out.txt'\n",
+    "    p.write_text(\"from pathlib import Path\\n\"\n",
+    "                 \"@call_parse\\n\"\n",
+    "                 f\"def inner(msg:str='none'): Path({str(out)!r}).write_text(msg)\")\n",
+    "    try:\n",
+    "        @call_parse(nested=True)\n",
+    "        def _runner(): runpy.run_path(str(p), init_globals={'call_parse':call_parse}, run_name='__main__')\n",
+    "        _runner()\n",
+    "    finally: sys.argv = argv\n",
+    "    test_eq(out.read_text(), 'hello')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0055556a",
+   "metadata": {},
+   "source": [
+    "`call_parse` auto-executes the function when its file is run directly, e.g `python foo.py` or `%run foo.py`. Importing a module never runs its CLI functions, and neither does defining one in a notebook cell:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "963c7a70",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ran=[]\n",
+    "@call_parse\n",
+    "def cell_cmd(): ran.append(1)\n",
+    "test_eq(ran, [])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8ec22c23",
+   "metadata": {},
+   "source": [
+    "This is the main way to use `fastcore.script`; decorate your function with `call_parse`, add `Param` annotations (as shown above) or type annotations and docments, and it can then be used as a script."
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "0d25da4b",
@@ -901,7 +1012,7 @@
   "solveit": {
    "default_code": false,
    "mode": "learning",
-   "use_thinking": false,
+   "use_thinking": true,
    "use_tools": true,
    "ver": 2
   }