Lua Import System for ComputerCraft

import.lua

#!/usr/bin/lua
--[[*
    * CCImport
    * (c) 2012 Orochimarufan
    *
    * a simple, python-style import mechanism.
    * made for ComputerCraft
    *
    * load with:
    * - require
    * - dofile
    * - os.loadAPI
    *
    * Will store itself in os.import, hence is only loaded once
    *
    * Use this line at the beginning of your file:
    * 'local import = os.import or require and require "import" or dofile "import.lua"'
    *
    * a module's metatable (_G metatable being inside the module) must not be messed with
    * export a __call__ function in your module to make it callable
    *
    * you can create importable standalone programs like in python
    * (the __name__ will be nil if you're the main script):
    *
    *   if __name__ == nil then
    *       main(...)
    *   end
    *
    * special variables:
    * import.path:          table:  where to search for modules ({"/"}) [py: sys.path]
    * import.code_ext:      string: module file extension (".lua")
    * import.package_code:  string: package __init__ filename ("__init__.lua")
    * import.modules:       table:  module table [py: sys.modules]
    * import.debug:         bool:   print some extra stuff (false)
    *
    * if imported file chunks return a table, it will be merged into the module's _G
    * that way, both the "returm M" and "Export Global Functions" approach works
    *
    * some people prefer to store their programs without the .lua extension,
    * you can set import.check_noext to try and import files without extension
    *
    * you could load this with CC's os.loadAPI() but it will leave a bogus import global
    * also, os.unloadAPI() will not unload it. use os.import=nil afterwards
    * using this with loadAPI only makes sense if you want to load it automatically (i.e /rom/apis)
    *
    * - Relative imports are supported, starting with '.'s (one for each level)
    * - 'local mod = import "pkg.mod"' is like 'from pkg import mod'
    * - 'local pkg = import.py "pkg.mod"' corresponds to 'import pkg.mod'
    * Warning: import.py ".lol" will give the TOP package, not necessarily "."
    * import.py "." will give the most toplevel package in your module tree
    *
    *]]

--[[
    TODO: move to io.*
    TODO: remove fs.*
  ]]

-- wether or not this will be called using os.loadAPI
-- false will work with import=dofile("import.lua")
-- it will store a copy of itself to os.import
-- if false, re-dofiling it will return os.import
-- CC will create an empty "import" table in _G.
-- as it is empty, it won't work. use os.import instead
-- api use is not recommended.
local api = false

-- api magic
if os.import then return os.import end
local import = {}
os.import = import

-- Settings
import.code_ext = ".lua"
import.package_code = "__init__.lua"
import.check_noext = false
import.debug = false

-- Import Path
local init = function()
if import._isCC then
    -- default path on CC
    import.path = {"/lib"}
else
    -- default path on PC
    import.path = {"."}
end
end

-- Module table
import._builtins = {
    -- common batteries
    string = string,
    table = table,
    coroutine = coroutine,
    os = os,
    io = io,
    bit = bit or bit32,
    math = math,
}
local CCBatteries = {
    -- CC batteries
    redstone = redstone,
    gps = gps,
    peripheral = peripheral,
    disk = disk,
    colors = colors,
    colours = colours,
    rednet = rednet,
    fs = fs,
    parallel = parallel,
    help = help,
    rs = rs,
    textutils = textutils,
    term = term,
    vector = vector,
}
local PCBatteries = {
    -- C Lua Batteries
    debug = debug,
    package = package,
}

-- initialize the modules table
import.modules = setmetatable({},{
    __tostring=function()return"<ModuleTable>"end,
    __index=import._builtins
})

-- Misc Helper Functions
local function strsplit(self, sep, max, is_regexp)
    assert(sep ~= '')
    assert(max == nil or max >= 1)

	local aRecord = {}

	if self:len() > 0 then
		local is_plain = not is_regexp
		local max = max or -1

		local field, start = 1, 1
		local first, last = self:find(sep, nil, is_plain)
		while first and max ~= 0 do
			aRecord[field] = self:sub(start, first-1)
			field, start, max = field+1, last+1, max-1
			first, last = self:find(sep, start, is_plain)
		end
		aRecord[field] = self:sub(start)
	end

	return aRecord
end

local function strjoin(delimiter, list)
  local len = #list
  if len == 0 then 
    return "" 
  end
  local string = list[1]
  for i = 2, len do 
    string = string .. delimiter .. list[i] 
  end
  return string
end

local function strstarts(str, piece)
    return string.sub(str, 1, string.len(piece)) == piece
end
local function strends(str, piece)
    return string.sub(str, -string.len(piece)) == piece
end

local function tblslice (values,i1,i2)
    local res = {}
    local n = #values
    -- default values for range
    i1 = i1 or 1
    i2 = i2 or n
    if i2 < 0 then
        i2 = n + i2 + 1
    elseif i2 > n then
        i2 = n
    end
    if i1 < 1 or i1 > n then return {} end
    local k = 1
    for i = i1,i2 do
        res[k] = values[i]
        k = k + 1
    end
    return res
end

-- C Lua 5.2
local setfenv = setfenv or function(f, t)
    f = (type(f) == 'function' and f or debug.getinfo(f + 1, 'f').func)
    local name
    local up = 0
    repeat
        up = up + 1
        name = debug.getupvalue(f, up)
    until name == '_ENV' or name == nil
    if name then
        debug.upvaluejoin(f, up, function() return t end, 1) -- use unique upvalue, set it to f
    end
end

local getfenv = getfenv or function(f)
    f = type(f)=='function' and f or debug.getinfo(f+1, 'f').func
    local name,up,env =nil,0,nil
    repeat
        up = up+1
        name,env = debug.getupvalue(f, up)
    until name=='_ENV' or name==nil
    return env
end

-- PC Lua
-- does not have the CC fs.* module
local fs = fs or {
    exists = function(_sFileName)
        local fp = io.open(_sFileName)
        if fp ~= nil then
            fp:close()
            --print(_sFileName..": OK")
            return true
        end
        --print(_sFileName..": not found")
    end,
    open = function(_sFileName, _sMode)
        assert(_sMode == "r", "only reading support atm. use io.open")
        local fp = io.open(_sFileName, _sMode)
        return {
            readAll = function()
                return fp:read("*a")
            end,
            readLine = function()
                return fp:read("*l")
            end,
            close = function()
                return fp:close()
            end }
    end,
    combine = function(_sBase, _sTail)
        if strends(_sBase, "/") then
            return _sBase.._sTail
        else
            return _sBase.."/".._sTail
       end
   end,
   getName = function(_sFileName)
       return strsplit(_sFileName, "/")[-1]
   end
}

-- export helpers as _import
-- this is not a part of the public api!
import._builtins._import = {
    setfenv = setfenv,
    getfenv = getfenv,
    strsplit = strsplit,
    strjoin = strjoin,
    strstarts = strstarts,
    strends = strends,
    tblslice = tblslice,
    fs = fs
}

-- check wether we're on CC or not
import._isCC = os.version and (strstarts(os.version(), "CraftOS") or strstarts(os.version(), "TurtleOS"))
init()
if import._isCC then
    -- CC batteries
    for bi, n in next,CCBatteries,nil do
        import._builtins[bi]=n
    end
else
    -- C Lua Batteries
    for bi, n in next,PCBatteries,nil do
        import._builtins[bi]=n
    end
end

-- Import ENV MT - Module Proxy
import.__envMT = {
    name_store = "__ModuleExecutionProxy__",
    __index = function(self, name)
        local module = rawget(self, "__module")[name]
        if module ~= nil then
            return module
        else
            return rawget(self, "__global")[name]
        end
    end,
    __newindex = function(self, name, value)
        rawget(self, "__module")[name] = value
    end,
    __tostring = function(self)
        return "<Module Execution Proxy: "..rawget(rawget(self, "__module"), "__name__")..">"
    end,
    ModEnv = function(module, global)
        local global    = global or _G
        local sStore    = import.__envMT.name_store
        local env       = rawget(module, sStore)
        if env then
            if rawget(env, "__global") ~= global then
                print("ImportWarning: Redefinition of "..tostring(env).." has different _G. Overriding!")
                rawset(env, "__global", global)
            end
        else
            env = setmetatable({}, import.__envMT)
            rawset(env, "__module", module)
            rawset(env, "__global", global)
            rawset(module, sStore, env)
        end
        return env
    end
}

-- Module MT
import.__modMT = {
    __tostring = function(self)
        if self.__isPkg__ then
            return "<Package: "..self.__name__..">"
        else
            return "<Module: "..self.__name__..">"
        end
    end,
    __newindex = function(self, name, value)
        if name == "__name__" or name == "__isPkg__" or name == "__fqn__" or
           name == "__path__" or name == import.__envMT.name_store then
           error("Cannot assign to reserved name "..name.." on module "..tostring(self), 2)
        end
        rawset(self, name, value)
    end,
    __call = function(self, ...)
        local fnc = rawget(self, "__call__")
        if not fnc then
            error("Module not callable: "..tostring(self), 2)
        end
        return fnc(...)
    end
}

-- Protected functions
function import._load(_sFile)
    -- have to handle shebang and such
	local file = fs.open(_sFile, "r")
	if file then
	    local firstline = file.readLine()
	    local data = file.readAll()
	    if not strstarts(firstline,"#!") then
	        data = firstline .. data end
		local func, err = loadstring(data, fs.getName( _sFile ))
		file.close()
		return func, err
	end
	return nil, "File not found"
end

--[[function import._exec_errh(o)
    local exc = {o}
    if debug then exc.traceback=debug.traceback(o, 2) end
    return exc
end]]
function import._exec_errh(o)
    if debug then return debug.traceback(o,2) end
    return o
end

function import._exec(co, env)
    -- execute module code in its table
    setfenv(co, env)
    local state, err = xpcall(co, import._exec_errh)
    if not state then return state, err end
    if type(err) == "table" then
        for name, thing in pairs(err) do
            env[name]=thing
        end
    end
    return state, env
end

function import._mk_comp(name)
    return strsplit(name, ".")
end

function import._get_module(path)
    if #path>0 then
        local current = import.modules[path[1]]
        if #path>1 then
            for i=2, #path do
                if current==nil then return nil end
                current = rawget(current, path[i])
            end
        end
        return current
    else
        return import.modules
    end
end

function import._have_module(path)
    return import._get_module(path) ~= nil
end

local function debugexists(path)
    if fs.exists(path) then
        print(path..": OK")
        return true
    else
        print(path..": Not found")
        return false
    end
end

function import._mk_fspath(lookin, path)
    local name = path[#path]
    local base = fs.combine(lookin, strjoin("/", tblslice(path, 1, #path-1)))
    -- generate paths
    local module = fs.combine(base, name .. import.code_ext)
    local package = fs.combine(base, name) .. "/" .. import.package_code
    -- return
    return package, module
end

function import._find_module(path, package)
    for _, lookin in ipairs(import.path) do
        local pkg, mod = import._mk_fspath(lookin, path)
        if fs.exists(pkg) then
            if fs.exists(mod) then
                print("ImportWarning: "..strjoin(".", path)..
                " has both a package and module file. ignoring module") end
            return pkg, true
        elseif package ~= true and fs.exists(mod) then
            return mod, false
        elseif package ~= true and import.check_noext then
            local file_noext = string.sub(mod,1,-import.code_ext:len())
            if fs.exists(file_noext) then
                return file_noext, nil
            end
        end
    end
    return nil, nil
end

function import._new_module(path, file, is_package)
    local module = {}
    module.__name__ = path[#path]
    module.__path__ = file
    module.__isPkg__ = is_package
    module.__fqn__ = strjoin(".", path)
    setmetatable(module, import.__modMT)
    import._insert_module(path, module)
    return module
end

function import._insert_module(path, module)
    local parent = tblslice(path, 1, -2)
    local name = path[#path]
    local pmodule = import._get_module(parent)
    assert(pmodule, "attempt to insert module to nonexistent parent")
    rawset(pmodule, name, module)
end

function import._clear_module(module)
    for n, v in pairs(module) do
        if str ~= "__name__" and str ~= "__path__" and str ~= "__fqn__"
            and str ~= "__isPkg__" and not strstarts(n, "__keep_") then
            module[n] = nil
        end
    end
end

function import._import_into(file, module)
    -- do not use on non-module-tables (that have a custom MT)
    -- load
    local co, err = import._load(file)
    if not co then return nil, err end
    -- environment proxy
    local env = import.__envMT.ModEnv(module, _G)
    -- execute
    return import._exec(co, env)
end

function import._import_module(path, package)
    local file, isPkg = import._find_module(path, package)
    if file == nil then
        error("ImportError: Module/Package not found: '"..strjoin(".",path).."'\n"..
              "        ".."import.path was: ['"..strjoin("', '", import.path).."']", 4) end
    local module = import._new_module(path, file, isPkg)
    if import.debug then print("Loading "..tostring(module).." from "..file) end
    local ok, err = import._import_into(file, module)
    if not ok then
        import._insert_module(path)
        error(err, 4)
    end
    return module
end

function import._relative(components)
    local levels = 1
    while components[levels+1]=="" do
        levels=levels+1
    end
    -- Lua 5.1 keeps tail calls on the stack, which is exactly what they are not supposed to be!?
    local ok, fenv = pcall(getfenv,4)
    if not ok then
        if fenv == "no function environment for tail call at level 4" then
            fenv = getfenv(4)
        else
            error(fenv)
        end
    end
    if not fenv.__fqn__ then
        error("RelativeImport: Caller does not seem to be located in a module (or has been setfenv()ed)", 3)
    end
    local parent = import._mk_comp(fenv.__fqn__)
    local comp = {}
    if levels==#components then
        levels = levels - 1
    end
    if levels > #parent then
        error(("RelativeImport: Cannot import up %s levels from %s"):format(levels,fenv.__fqn__), 3)
    end
    for i=1,#parent-levels do
        table.insert(comp, parent[i])
    end
    if components[levels+1]~="" then
        for i=levels+1,#components do
            table.insert(comp, components[i])
        end
    end
    return comp
end

-- main importer
function import._import(components)
    -- no empty fqn components
    for _, comp in next,components,nil do
        if comp=="" then
            error("Empty Component in FQN: "..strjoin(".",components).." (#".._..")", 3)
        end
    end
    -- look if we already imported it
    local module = import._get_module(components)
    if module then
        if import.debug then
            print(("Have %s from %s in module table"):format(module,module.__path__))
        end
        return module
    end
    -- import parents
    for i=1, #components-1 do
        local path = tblslice(components, 1, i)
        if not import._have_module(path) then
            import._import_module(path, true)
        end
    end
    -- import module
    return import._import_module(components)
end

-- Public
function import.import(name)
    local path = import._mk_comp(name)
    -- relative
    if path[1] == "" then
        path = import._relative(path)
    end
    -- import
    return import._import(path)
end

function import.py(name)
    local path = import._mk_comp(name)
    if path[1] == "" then
        path = import._relative(path)
    end
    import._import(path)
    return import._get_module(import.modules[path[1]])
end

function import.clearm0d5()
    -- clear module cache import.modules
    -- existing references to modules will persist!
    for n,m in pairs(import.modules) do
        import.modules[n] = nil
    end
end

--[[function import.reload(name)
    local components = import._mk_comp(name)
    assert(import._have_module(components), name.." was never imported")
    return import._reload_module(path)
end ]]--

-- metatable magic
import.__name__ = "import"
import.__isPkg__ = false
import.__path__ = nil
import.__fqn__ = "import"
setmetatable(import, import.__modMT)

-- make the module callable
import.__call__ = import.import

return import