cli.py 6.1 KB

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