conda_env_path_add.py

#!/usr/bin/env python

import os, sys, subprocess

def main():
	conda_base_path = os.path.dirname(os.path.dirname(__file__))
	path = sys.argv[1]
	env = get_env()
	gen_activate_script(path, env, conda_base_path)


def get_env():
	if len(sys.argv) > 2:
		env = sys.argv[2]
	else:
		print("Scanning conda for available environments...")
		env_list = subprocess.check_output(["conda","env","list"]).decode("utf-8")
		print(env_list)
		while True:
			env = input("Please enter the name of the conda environment you would like to change: ").strip()
			if check_valid_env(env_list, env) == True:
				break
	return env


def check_valid_env(env_list, env):
	for line in env_list.split('\n'):
		if line.startswith('#'):
			continue
		if env == line.split()[0]:
			return True
	return False


def gen_activate_script(path, env, conda_base_path):
	activate_dir_path = os.path.join(conda_base_path,"envs",env,"etc","conda","activate.d")
	os.makedirs(activate_dir_path, exist_ok=True)
	if os.path.isfile(os.path.join(activate_dir_path,"env_vars.sh")):
		with open(os.path.join(activate_dir_path,"env_vars.sh")) as activate_handle:
			with open(os.path.join(activate_dir_path,"env_vars_temp.sh"),"w") as activate_temp_handle:
				for line in activate_handle:
					if line.startswith("export PATH="):
						if path in line.strip().split(":"):
							print("{} has already been added to the PATH variable of environment {}. Exiting...".format(path, env))
							activate_temp_handle.write(line)
						else:
							activate_temp_handle.write("{}:{}\n".format(line.strip(), path))
							print("$PATH variable in {} has been successfully updated.".format(env))
					else:
						activate_temp_handle.write(line)
		os.remove(os.path.join(activate_dir_path,"env_vars.sh"))
		os.rename(os.path.join(activate_dir_path,"env_vars_temp.sh"),os.path.join(activate_dir_path,"env_vars.sh"))
	else:
		with open(os.path.join(activate_dir_path,"env_vars.sh"),"w") as activate_handle:
			activate_handle.write("export PATH=$PATH:{}\n".format(path))


main()