"main" function in Lua?

Question

In python, one would usually define a main function, in order to allow the script to be used as module (if needed):

def main():
    print("Hello world")
    return 0

if __name__ == "__main__":
    sys.exit(main())

In Lua, the idiom if __name__ == "__main__" isn't possible as such (that means, I don't think it is).

That's what I'm usually doing in order to have a similar behaviour in Lua:

os.exit((function(args)
    print("Hello world")
    return 0
end)(arg))

... But this approach seems rather "heavy on parentheses" :-)

Is there a more common approach (besides defining a global main function, which seems redundant)?

Solution 1

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

Solution 2

When Lua requires a module, it passes it the name it's been required with as varargs (...).

So, if your script doesn't intend to take any arguments (from the command line or otherwise), you can use something like

if ... then
  return this_mod --module case
else
  main() --main case
end

Note, however, that this isn't foolproof in the (entirely) possible case that you take arguments. However, at this point, you can combine this with Lukasz's answer to get:

if not package.loaded[...] then
  --main case
else --module case
end

Still not perfect (for instance, if the script is called with a first argument of string or the name of some other already-loaded module), but likely good enough. In other situations, I defer to MattJ's answer.

Solution 3

you could try to check if the module has been required.

from documentation:

package.loaded A table used by require to control which modules are already loaded. When you require a module modname and package.loaded[modname] is not false, require simply returns the value stored there.

With this you could write:

if not package.loaded['modulename'] then
    main()
end

Solution 4

What's wrong with this:

$ cat aa.lua
#!/usr/bin/lua

if (arg ~= nil and arg[-1] ~= nil) then
    print "main"
else
    print "library"
end
$ ./aa.lua
main
$ ./aa.lua arg1 arg2
main
$ cat bb.lua
#!/usr/bin/lua

print("in bb")
$ lua -laa bb.lua
library
in bb
$ lua
Lua 5.1.4  Copyright (C) 1994-2008 Lua.org, PUC-Rio
> require "aa"
library
> 

Solution 5

I am going to suggest yet another variation, that seems to work on lua5.1 as well as lua5.2:

function is_main(offset)
    return debug.getinfo(4 + (offset or 0)) == nil
end

if is_main() then
    print("Main chunk!")
else
    print("Library chunk!")
end

If you do not feel like defining an extra function is_main for this purpose, you can just do:

if debug.getinfo(3) == nil then
    print("Main chunk!")
else
    print("Library chunk!")
end

View more solutions

Comments

[Admin]

downside is, that you'd have to know the name of the module... arg[0] wouldn't work either, as lua modules are matched after ?.lua;?.BIN (BIN being the file extension for a native library, .so or .dll for example)

[Łukasz Gruner]

what are you trying to accomplish? I assumed that you want to put this code into your own module, and know how it will be named.

[Łukasz Gruner]

btw, arg[0] returns the name of lua file that is being run, or that has required your module. try: test1.lua > require "test" test.lua > print(arg[0])

[Admin]

"btw, arg[0] returns the name of lua file that is being run" - i've already mentioned that in my previous comment...

[Admin]

I wish I could upvote more than once :-). I guess i'll stick with the getfenv method for now, as it seems smaller. i'll definitely keep the other solution in mind.

[coldfix]

The problem is arg could be defined as a global variable in the calling code. Otherwise this is simple and great!

[coldfix]

Great idea. Unfortunately it does not work, since module code is considered 'main' chunk too.

[coldfix]

This does not work in lua5.2. The package.loaded entry will be inserted only when the module returns.

[coldfix]

The second suggestion does not work in lua5.2. The package.loaded entry will be inserted only when the module returns

[coldfix]

One should mention that this has a similar downfall as MattJ's answer: If arg is defined as a global variable and contains exactly the module name, you will think that the module is called as a main file.

[xjshiya]

@coldfix I've added some modified code that seems to work with 5.2. The rest of my answer still stands though, use at your own risk :)

[coldfix]

This looks good and much cleaner/less error prone than the second method. Is the corresponding scope guaranteed to have at least one entry? I posted a very similar solution using debug.getinfo instead of debug.getlocal, which makes the use of pcall unnecessary. Regards, Thomas

[xjshiya]

Yes, arguably debug.getinfo() was the correct approach from the start. However I guess I was originally trying to avoid depending on the debug library if I didn't have to...

[selurvedu]

Thanks, the updated answer works in Lua 5.1, 5.2 and 5.3.

[user3479901]

This doesn't work if more than one module in the dependency tree does this: at best, you'd have to have a "FOOBAR_CLI=1" (where "foobar" is the name of the module) for each module in the project - and forget about having more than one instance of that module (eg. to resolve a dependency version conflict)!