“main” function in Lua?

There’s no “proper” way to do this, since Lua doesn’t really distinguish code by where it came from, they are all just functions. That said, this at least seems to work in Lua 5.1:

matthew@silver:~$ cat hybrid.lua 
if pcall(getfenv, 4) then
    print("Library")
else
    print("Main file")
end
matthew@silver:~$ lua hybrid.lua 
Main file
matthew@silver:~$ lua -lhybrid
Library
Lua 5.1.4  Copyright (C) 1994-2008 Lua.org, PUC-Rio
> ^C
matthew@silver:~$ lua
Lua 5.1.4  Copyright (C) 1994-2008 Lua.org, PUC-Rio
> require "hybrid"
Library
> ^C
matthew@silver:~$

It works by checking whether the stack depth is greater than 3 (the normal depth for a file in the stock Lua interpreter). This test may break between Lua versions though, and even in any embedded/custom Lua builds.

I’ll also include this (slightly more portable) alternative, although it’s taking an even greater leap in heuristics, and has a failure case (see below):

matthew@silver:~$ cat hybrid2.lua 
function is_main(_arg, ...)
    local n_arg = _arg and #_arg or 0;
    if n_arg == select("#", ...) then
        for i=1,n_arg do
            if _arg[i] ~= select(i, ...) then
                print(_arg[i], "does not match", (select(i, ...)))
                return false;
            end
        end
        return true;
    end
    return false;
end

if is_main(arg, ...) then
    print("Main file");
else
    print("Library");
end
matthew@silver:~$ lua hybrid2.lua 
Main file
matthew@silver:~$ lua -lhybrid2
Library
Lua 5.1.4  Copyright (C) 1994-2008 Lua.org, PUC-Rio
> ^C
matthew@silver:~$ lua 
Lua 5.1.4  Copyright (C) 1994-2008 Lua.org, PUC-Rio
> require "hybrid2"
Library
>

This one works by comparing the contents of _G.arg with the contents of ‘…’. In the main chunk they will always be the same. In a module _G.arg will still contain the command-line arguments, but ‘…’ will contain the module name passed to require(). I suspect this is closer to the better solution for you, given that you know your module name. The bug in this code lies when the user executes the main script with 1 argument, and this is the exact name of your module:

matthew@silver:~$ lua -i hybrid2.lua hybrid2
Lua 5.1.4  Copyright (C) 1994-2008 Lua.org, PUC-Rio
Main file
> require "hybrid2"
Main file
> 

Given the above I hope at least you know where you stand, even if it isn’t exactly what you had in mind 🙂

Update: For a version of hybrid.lua that works in Lua 5.1 and 5.2, you can replace getfenv with debug.getlocal:

if pcall(debug.getlocal, 4, 1) then
    print("Library")
else
    print("Main file")
end

Leave a Comment