environment.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. from __future__ import absolute_import
  2. from __future__ import unicode_literals
  3. import codecs
  4. import contextlib
  5. import logging
  6. import os
  7. import re
  8. import six
  9. from ..const import IS_WINDOWS_PLATFORM
  10. from .errors import ConfigurationError
  11. from .errors import EnvFileNotFound
  12. log = logging.getLogger(__name__)
  13. def split_env(env):
  14. if isinstance(env, six.binary_type):
  15. env = env.decode('utf-8', 'replace')
  16. key = value = None
  17. if '=' in env:
  18. key, value = env.split('=', 1)
  19. else:
  20. key = env
  21. if re.search(r'\s', key):
  22. raise ConfigurationError(
  23. "environment variable name '{}' may not contain whitespace.".format(key)
  24. )
  25. return key, value
  26. def env_vars_from_file(filename):
  27. """
  28. Read in a line delimited file of environment variables.
  29. """
  30. if not os.path.exists(filename):
  31. raise EnvFileNotFound("Couldn't find env file: {}".format(filename))
  32. elif not os.path.isfile(filename):
  33. raise EnvFileNotFound("{} is not a file.".format(filename))
  34. env = {}
  35. with contextlib.closing(codecs.open(filename, 'r', 'utf-8-sig')) as fileobj:
  36. for line in fileobj:
  37. line = line.strip()
  38. if line and not line.startswith('#'):
  39. try:
  40. k, v = split_env(line)
  41. env[k] = v
  42. except ConfigurationError as e:
  43. raise ConfigurationError('In file {}: {}'.format(filename, e.msg))
  44. return env
  45. class Environment(dict):
  46. def __init__(self, *args, **kwargs):
  47. super(Environment, self).__init__(*args, **kwargs)
  48. self.missing_keys = []
  49. self.silent = False
  50. @classmethod
  51. def from_env_file(cls, base_dir, env_file=None):
  52. def _initialize():
  53. result = cls()
  54. if base_dir is None:
  55. return result
  56. if env_file:
  57. env_file_path = os.path.join(base_dir, env_file)
  58. else:
  59. env_file_path = os.path.join(base_dir, '.env')
  60. try:
  61. return cls(env_vars_from_file(env_file_path))
  62. except EnvFileNotFound:
  63. pass
  64. return result
  65. instance = _initialize()
  66. instance.update(os.environ)
  67. return instance
  68. @classmethod
  69. def from_command_line(cls, parsed_env_opts):
  70. result = cls()
  71. for k, v in parsed_env_opts.items():
  72. # Values from the command line take priority, unless they're unset
  73. # in which case they take the value from the system's environment
  74. if v is None and k in os.environ:
  75. result[k] = os.environ[k]
  76. else:
  77. result[k] = v
  78. return result
  79. def __getitem__(self, key):
  80. try:
  81. return super(Environment, self).__getitem__(key)
  82. except KeyError:
  83. if IS_WINDOWS_PLATFORM:
  84. try:
  85. return super(Environment, self).__getitem__(key.upper())
  86. except KeyError:
  87. pass
  88. if not self.silent and key not in self.missing_keys:
  89. log.warning(
  90. "The {} variable is not set. Defaulting to a blank string."
  91. .format(key)
  92. )
  93. self.missing_keys.append(key)
  94. return ""
  95. def __contains__(self, key):
  96. result = super(Environment, self).__contains__(key)
  97. if IS_WINDOWS_PLATFORM:
  98. return (
  99. result or super(Environment, self).__contains__(key.upper())
  100. )
  101. return result
  102. def get(self, key, *args, **kwargs):
  103. if IS_WINDOWS_PLATFORM:
  104. return super(Environment, self).get(
  105. key,
  106. super(Environment, self).get(key.upper(), *args, **kwargs)
  107. )
  108. return super(Environment, self).get(key, *args, **kwargs)
  109. def get_boolean(self, key):
  110. # Convert a value to a boolean using "common sense" rules.
  111. # Unset, empty, "0" and "false" (i-case) yield False.
  112. # All other values yield True.
  113. value = self.get(key)
  114. if not value:
  115. return False
  116. if value.lower() in ['0', 'false']:
  117. return False
  118. return True