import asyncio
import os
import re
import sys

from collections import OrderedDict
from functools import partial, wraps

LMOD_SYSTEM_NAME = os.environ.get("LMOD_SYSTEM_NAME", "")
SITE_POSTFIX = os.path.join("lib", "python" + sys.version[:3], "site-packages")

MODULE_REGEX = re.compile(r"^([a-zA-z0-9-_+.]{1,}\/[a-zA-z0-9-_+.]{1,})$", re.M)
MODULE_REGEX_NO_HIDDEN =  re.compile(r"^([a-zA-z0-9-_+.]{1,}\/[a-zA-z0-9-_][a-zA-z0-9-_.]*)$", re.M)

async def module(command, *args):
    cmd = os.environ["LMOD_CMD"], "python", "--terse", command, *args

    proc = await asyncio.create_subprocess_shell(
        " ".join(cmd), stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
    )

    stdout, stderr = await proc.communicate()

    if command in ("load", "unload", "purge", "restore", "save"):
        try:
            exec(stdout.decode())
        except NameError:
            pass

    if stderr:
        return stderr.decode()


def update_sys_path(env_var, postfix=""):
    def decorator(function):
        @wraps(function)
        def wrapper(*args, **kargs):
            if env_var in os.environ:
                orig_paths = os.environ[env_var].split(os.pathsep)
                orig_paths = [os.path.join(path, postfix) for path in orig_paths]
            else:
                orig_paths = []
            output = function(*args, **kargs)
            if env_var in os.environ:
                new_paths = os.environ[env_var].split(os.pathsep)
                new_paths = [os.path.join(path, postfix) for path in new_paths]
            else:
                new_paths = []

            paths_to_del = set(orig_paths) - set(new_paths)
            paths_to_add = set(new_paths) - set(orig_paths)
            paths_to_add = [path for path in new_paths if path in paths_to_add]
            for path in paths_to_del:
                try:
                    sys.path.remove(path)
                except ValueError:
                    continue
            sys.path.extend(paths_to_add)
            return output

        return wrapper

    return decorator


class API(object):
    def __init__(self, show_cache_capacity=128):
        self.avail_cache = None
        self.list_cache = None
        self.savelist_cache = None
        self.show_cache = OrderedDict()
        self.show_cache_capacity = show_cache_capacity

    def invalidate_module_caches(self):
        self.avail_cache = None
        self.list_cache = None
        self.show_cache = OrderedDict()

    async def avail(self, *args):
        if self.avail_cache is None:
            string = await module("avail", *args)
            modules = re.findall(MODULE_REGEX, string)
            modules.sort(key=lambda v: v.split("/")[0])
            self.avail_cache = modules
        return self.avail_cache

    async def list(self, include_hidden=False):
        if self.list_cache is None:
            string = await module("list")
            string = string.strip()
            if string != "No modules loaded":
                regex = MODULE_REGEX_NO_HIDDEN if not include_hidden else MODULE_REGEX
                self.list_cache = re.findall(regex, string)
            else:
                self.list_cache = []
        return self.list_cache

    async def freeze(self):
        modules = await self.list(include_hidden=False)
        return "\n".join(
            [
                "import lmod",
                "await lmod.purge(force=True)",
                "await lmod.load({})".format(str(modules)[1:-1]),
            ]
        )

    @update_sys_path("PYTHONPATH")
    @update_sys_path("EBPYTHONPREFIXES", SITE_POSTFIX)
    async def load(self, *modules):
        output = await module("load", *modules)
        if output:
            print(output)
        self.invalidate_module_caches()

    @update_sys_path("PYTHONPATH")
    @update_sys_path("EBPYTHONPREFIXES", SITE_POSTFIX)
    async def reset(self):
        output = await module("reset")
        if output:
            print(output)
        self.invalidate_module_caches()

    @update_sys_path("PYTHONPATH")
    @update_sys_path("EBPYTHONPREFIXES", SITE_POSTFIX)
    async def restore(self, name):
        output = await module("restore", name)
        if output:
            print(output)
        self.invalidate_module_caches()

    async def save(self, name):
        output = await module("save", name)
        if output:
            print(output)
        self.savelist_cache = None

    async def savelist(self):
        if self.savelist_cache is None:
            string = await module("savelist")
            self.savelist_cache = string.split()
        return self.savelist_cache

    @update_sys_path("PYTHONPATH")
    @update_sys_path("EBPYTHONPREFIXES", SITE_POSTFIX)
    async def unload(self, *modules):
        output = await module("unload", *modules)
        if output:
            print(output)
        self.invalidate_module_caches()

    @update_sys_path("PYTHONPATH")
    @update_sys_path("EBPYTHONPREFIXES", SITE_POSTFIX)
    async def purge(self, force=False):
        if force:
            args = ("--force",)
        else:
            args = ()
        output = await module("purge", *args)
        if output:
            print(output)
        self.invalidate_module_caches()

    async def show(self, name):
        if not name in self.show_cache:
            result = await module("show", name)
            self.show_cache[name] = result
            self.show_cache.move_to_end(name)
            if len(self.show_cache) > self.show_cache_capacity:
                self.show_cache.popitem(last = False)
        return self.show_cache[name]


_lmod = API()

avail = _lmod.avail
list = _lmod.list
freeze = _lmod.freeze
load = _lmod.load
reset = _lmod.reset
restore = _lmod.restore
save = _lmod.save
savelist = _lmod.savelist
unload = _lmod.unload
purge = _lmod.purge
show = _lmod.show

@@ -40,16 +40,96 @@
 rt sys%0A%0A
+from asyncio import create_subprocess_shell%0Afrom asyncio.subprocess import PIPE%0A
 from col
@@ -194,16 +194,50 @@
  wraps%0A%0A
+LMOD_CMD = os.environ%5B%22LMOD_CMD%22%5D%0A
 LMOD_SYS
@@ -594,36 +594,24 @@
   cmd = 
-os.environ%5B%22
 LMOD_CMD
 %22%5D, %22pyt
@@ -602,18 +602,16 @@
 LMOD_CMD
-%22%5D
 , %22pytho
@@ -658,24 +658,16 @@
 = await 
-asyncio.
 create_s
@@ -717,59 +717,21 @@
 out=
-asyncio.subprocess.PIPE, stderr=asyncio.subprocess.
+PIPE, stderr=
 PIPE
@@ -830,16 +830,25 @@
 %22purge%22,
+ %22reset%22,
  %22restor

Clean module function