cli.py 5.7 KB

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