cli.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  1. import json
  2. import os
  3. import shlex
  4. import sys
  5. from contextlib import contextmanager
  6. from typing import IO, Any, Dict, Iterator, List, Optional
  7. if sys.platform == "win32":
  8. from subprocess import Popen
  9. try:
  10. import click
  11. except ImportError:
  12. sys.stderr.write(
  13. "It seems python-dotenv is not installed with cli option. \n"
  14. 'Run pip install "python-dotenv[cli]" to fix this.'
  15. )
  16. sys.exit(1)
  17. from .main import dotenv_values, set_key, unset_key
  18. from .version import __version__
  19. def enumerate_env() -> Optional[str]:
  20. """
  21. Return a path for the ${pwd}/.env file.
  22. If pwd does not exist, return None.
  23. """
  24. try:
  25. cwd = os.getcwd()
  26. except FileNotFoundError:
  27. return None
  28. path = os.path.join(cwd, ".env")
  29. return path
  30. @click.group()
  31. @click.option(
  32. "-f",
  33. "--file",
  34. default=enumerate_env(),
  35. type=click.Path(file_okay=True),
  36. help="Location of the .env file, defaults to .env file in current working directory.",
  37. )
  38. @click.option(
  39. "-q",
  40. "--quote",
  41. default="always",
  42. type=click.Choice(["always", "never", "auto"]),
  43. help="Whether to quote or not the variable values. Default mode is always. This does not affect parsing.",
  44. )
  45. @click.option(
  46. "-e",
  47. "--export",
  48. default=False,
  49. type=click.BOOL,
  50. help="Whether to write the dot file as an executable bash script.",
  51. )
  52. @click.version_option(version=__version__)
  53. @click.pass_context
  54. def cli(ctx: click.Context, file: Any, quote: Any, export: Any) -> None:
  55. """This script is used to set, get or unset values from a .env file."""
  56. ctx.obj = {"QUOTE": quote, "EXPORT": export, "FILE": file}
  57. @contextmanager
  58. def stream_file(path: os.PathLike) -> Iterator[IO[str]]:
  59. """
  60. Open a file and yield the corresponding (decoded) stream.
  61. Exits with error code 2 if the file cannot be opened.
  62. """
  63. try:
  64. with open(path) as stream:
  65. yield stream
  66. except OSError as exc:
  67. print(f"Error opening env file: {exc}", file=sys.stderr)
  68. sys.exit(2)
  69. @cli.command(name="list")
  70. @click.pass_context
  71. @click.option(
  72. "--format",
  73. "output_format",
  74. default="simple",
  75. type=click.Choice(["simple", "json", "shell", "export"]),
  76. help="The format in which to display the list. Default format is simple, "
  77. "which displays name=value without quotes.",
  78. )
  79. def list_values(ctx: click.Context, output_format: str) -> None:
  80. """Display all the stored key/value."""
  81. file = ctx.obj["FILE"]
  82. with stream_file(file) as stream:
  83. values = dotenv_values(stream=stream)
  84. if output_format == "json":
  85. click.echo(json.dumps(values, indent=2, sort_keys=True))
  86. else:
  87. prefix = "export " if output_format == "export" else ""
  88. for k in sorted(values):
  89. v = values[k]
  90. if v is not None:
  91. if output_format in ("export", "shell"):
  92. v = shlex.quote(v)
  93. click.echo(f"{prefix}{k}={v}")
  94. @cli.command(name="set")
  95. @click.pass_context
  96. @click.argument("key", required=True)
  97. @click.argument("value", required=True)
  98. def set_value(ctx: click.Context, key: Any, value: Any) -> None:
  99. """
  100. Store the given key/value.
  101. This doesn't follow symlinks, to avoid accidentally modifying a file at a
  102. potentially untrusted path.
  103. """
  104. file = ctx.obj["FILE"]
  105. quote = ctx.obj["QUOTE"]
  106. export = ctx.obj["EXPORT"]
  107. success, key, value = set_key(file, key, value, quote, export)
  108. if success:
  109. click.echo(f"{key}={value}")
  110. else:
  111. sys.exit(1)
  112. @cli.command()
  113. @click.pass_context
  114. @click.argument("key", required=True)
  115. def get(ctx: click.Context, key: Any) -> None:
  116. """Retrieve the value for the given key."""
  117. file = ctx.obj["FILE"]
  118. with stream_file(file) as stream:
  119. values = dotenv_values(stream=stream)
  120. stored_value = values.get(key)
  121. if stored_value:
  122. click.echo(stored_value)
  123. else:
  124. sys.exit(1)
  125. @cli.command()
  126. @click.pass_context
  127. @click.argument("key", required=True)
  128. def unset(ctx: click.Context, key: Any) -> None:
  129. """
  130. Removes the given key.
  131. This doesn't follow symlinks, to avoid accidentally modifying a file at a
  132. potentially untrusted path.
  133. """
  134. file = ctx.obj["FILE"]
  135. quote = ctx.obj["QUOTE"]
  136. success, key = unset_key(file, key, quote)
  137. if success:
  138. click.echo(f"Successfully removed {key}")
  139. else:
  140. sys.exit(1)
  141. @cli.command(
  142. context_settings={
  143. "allow_extra_args": True,
  144. "allow_interspersed_args": False,
  145. "ignore_unknown_options": True,
  146. }
  147. )
  148. @click.pass_context
  149. @click.option(
  150. "--override/--no-override",
  151. default=True,
  152. help="Override variables from the environment file with those from the .env file.",
  153. )
  154. @click.argument("commandline", nargs=-1, type=click.UNPROCESSED)
  155. def run(ctx: click.Context, override: bool, commandline: tuple[str, ...]) -> None:
  156. """Run command with environment variables present."""
  157. file = ctx.obj["FILE"]
  158. if not os.path.isfile(file):
  159. raise click.BadParameter(
  160. f"Invalid value for '-f' \"{file}\" does not exist.", ctx=ctx
  161. )
  162. dotenv_as_dict = {
  163. k: v
  164. for (k, v) in dotenv_values(file).items()
  165. if v is not None and (override or k not in os.environ)
  166. }
  167. if not commandline:
  168. click.echo("No command given.")
  169. sys.exit(1)
  170. run_command([*commandline, *ctx.args], dotenv_as_dict)
  171. def run_command(command: List[str], env: Dict[str, str]) -> None:
  172. """Replace the current process with the specified command.
  173. Replaces the current process with the specified command and the variables from `env`
  174. added in the current environment variables.
  175. Parameters
  176. ----------
  177. command: List[str]
  178. The command and it's parameters
  179. env: Dict
  180. The additional environment variables
  181. Returns
  182. -------
  183. None
  184. This function does not return any value. It replaces the current process with the new one.
  185. """
  186. # copy the current environment variables and add the vales from
  187. # `env`
  188. cmd_env = os.environ.copy()
  189. cmd_env.update(env)
  190. if sys.platform == "win32":
  191. # execvpe on Windows returns control immediately
  192. # rather than once the command has finished.
  193. p = Popen(command, universal_newlines=True, bufsize=0, shell=False, env=cmd_env)
  194. _, _ = p.communicate()
  195. sys.exit(p.returncode)
  196. else:
  197. os.execvpe(command[0], args=command, env=cmd_env)