text: """This module aims to make it easier to run the Alembic migration
  scripts of Kotti and Kotti add-ons by providing a uniform access.
  
  Commands herein will typically be called by the console script
  ``kotti-migrate`` (see the docstring of that command below).
  
  Kotti stores the current revision of its migration in table
  ``kotti_alembic_versions``.  The convention here is
  ``<packagename>_alembic_versions``.  You should normally not need to
  worry about the name of this table, as it is created and managed
  automatically through this module.  If, however, you plan to use your
  own alembic.ini configuration file for your add-on or application,
  keep in mind to configure a table name as described above.  The table
  name can be set using Alembic's ``version_table`` option.
  
  Kotti has start-up code that will create the database from scratch if
  it doesn't exist.  This code will also call this module's function
  ``stamp_heads`` to set the current revision of all migrations
  registered with this module to the latest.  This assumes that when we
  create the database from scratch (using ``metadata.create_all``), we
  don't need to run any of the past migrations.
  
  Unfortunately, this won't help in the situation where a user adds an
  add-on with migrations to the Kotti site _after_ the database was
  initialized for the first time.  In this case, users of the add-on
  will need to run ``kotti-migrate stamp_head
  --scripts=yourpackage:alembic``, or the add-on author will have to
  write equivalent code somewhere in their populate hook.
  
  Add-on authors can register their Alembic scripts with this module by
  adding their Alembic 'script directory' location to the
  ``kotti.alembic_dirs`` setting.  An example::
  
    def kotti_configure(settings):
        # ...
        settings['kotti.alembic_dirs'] += ' kotti_contactform:alembic'
  
  ``kotti-migrate`` commands 'list_all', 'upgrade_all' and
  'stamp_heads' will then include the add-on.
  """
  import os
  from typing import Callable
  from typing import List
  
  import pkg_resources
  from alembic.config import Config
  from alembic.environment import EnvironmentContext
  from alembic.script import ScriptDirectory
  from alembic.util import load_python_file
  from zope.sqlalchemy import mark_changed
  
  from kotti import DBSession
  from kotti import conf_defaults
  from kotti import get_settings
  from kotti.util import command
  
  KOTTI_SCRIPT_DIR = pkg_resources.resource_filename("kotti", "alembic")
  DEFAULT_LOCATION = "kotti:alembic"
  
  
  class ScriptDirectoryWithDefaultEnvPy(ScriptDirectory):
      @property
      def env_py_location(self) -> str:
          loc = super().env_py_location
          if not os.path.exists(loc):
              loc = os.path.join(KOTTI_SCRIPT_DIR, "env.py")
          return loc
  
      def run_env(self) -> None:
          dir_, filename = self.env_py_location.rsplit(os.path.sep, 1)
          load_python_file(dir_, filename)
  
  
  class PackageEnvironment:
      def __init__(self, location: str) -> None:
          self.location = location
          self.config = self._make_config(location)
          self.script_dir = self._make_script_dir(self.config)
  
      @property
      def pkg_name(self) -> str:
          return self.location.split(":")[0]
  
      @property
      def version_table(self) -> str:
          return f"{self.pkg_name}_alembic_version"
  
      def run_env(self, fn: Callable, **kw) -> None:
          with EnvironmentContext(
              self.config, self.script_dir, fn=fn, version_table=self.version_table, **kw
          ):
              self.script_dir.run_env()
  
      @staticmethod
      def _make_config(location: str) -> Config:
          cfg = Config()
          cfg.set_main_option("script_location", location)
          cfg.set_main_option("sqlalchemy.url", get_settings()["sqlalchemy.url"])
          return cfg
  
      @staticmethod
      def _make_script_dir(alembic_cfg: Config) -> ScriptDirectoryWithDefaultEnvPy:
          script_dir = ScriptDirectory.from_config(alembic_cfg)
          script_dir.__class__ = ScriptDirectoryWithDefaultEnvPy  # O_o
          return script_dir
  
  
  def get_locations() -> List[str]:
      conf_str = get_settings()["kotti.alembic_dirs"]
      return [line.strip() for line in conf_str.split() if line.strip()]
  
  
  def stamp_head(location: str = DEFAULT_LOCATION, revision: None = None) -> None:
  
      env = PackageEnvironment(location)
  
      def do_stamp(rev, context, revision=revision):
  
          if revision is None:
              revision = context.script.get_current_head()
          elif revision == "None":
              revision = None
  
          context.stamp(env.script_dir, revision)
  
          mark_changed(DBSession())
          return []
  
      env.run_env(do_stamp)
  
  
  def stamp_heads() -> None:
      for location in get_locations():
          stamp_head(location)
  
  
  def upgrade(location=DEFAULT_LOCATION, revision=None):
      # We don't want to fire any kind of events during a migration,
      # because "migrations are a low-level thing".
      from kotti import events
  
      events.clear()
  
      pkg_env = PackageEnvironment(location)
  
      if revision is None:
          revision = pkg_env.script_dir.get_current_head()
      print(f"Upgrading {pkg_env.location}:")
  
      def upgrade(heads, context):
          # alembic supports multiple heads, we don't.
          # initial revision is () in alembic >= 0.7
          rev = heads[0] if heads else None
  
          if rev == revision:
              print("  - already up to date.")
              return []
  
          print(f"  - upgrading from {rev} to {revision}...")
  
          return context.script._upgrade_revs(revision, rev)
  
      pkg_env.run_env(upgrade, starting_rev=None, destination_rev=revision)
      print()
  
  
  def upgrade_all():
      for location in get_locations():
          upgrade(location)
  
  
  def list_all():
      pkg_envs = [PackageEnvironment(location) for location in get_locations()]
      for pkg_env in pkg_envs:
          print(f"{pkg_env.pkg_name}:")
  
          for script in pkg_env.script_dir.walk_revisions():
              print(
                  "  - {} -> {}: {}".format(
                      script.down_revision, script.revision, script.doc
                  )
              )
  
          def current_revision(rev, context):
              rev = rev[0] if rev else None
              print(f"  - current revision: {rev}")
              return []
  
          pkg_env.run_env(current_revision)
          print()
  
  
  def kotti_migrate_command():
      __doc__ = """Migrate Kotti and Kotti add-ons.
  
      Usage:
        kotti-migrate <config_uri> list_all
        kotti-migrate <config_uri> upgrade [--scripts=<location>] [--rev=<rev>]
        kotti-migrate <config_uri> upgrade_all
        kotti-migrate <config_uri> stamp_head [--scripts=<location>] [--rev=<rev>]
  
      o 'list_all' prints a list of all available migrations of Kotti
        and registered add-ons.
  
      o 'upgrade' will run Kotti's upgrades to upgrade the database to
      the latest version.
  
        Use '--scripts=kotti_myaddon:alembic' to run the upgrades of the
        'kotti_myaddon' package instead.
  
      o 'upgrade_all' will run all upgrades of all packages registered
        with Kotti.
  
      o 'stamp_head' allows you to manually set the stamped version to
        the latest version inside the 'kotti_alembic_version' table,
        that is, without actually running any migrations.
  
        You may use this command for a different package by using the
        '--scripts' option.
  
      Options:
        -h --help     Show this screen.
      """
      # We need to turn off populators and root_factory when we run
      # migrations, because they would access the database, which may
      # not be possible prior to the migration.
      #
      # Unfortunately, we're not able to just set the 'kotti.populators'
      # setting to an empty list.  Since add-ons might add to this list
      # again later, when we call 'bootstrap' (and thus their
      # 'includeme' function).
      save_conf_defaults = conf_defaults.copy()
  
      os.environ["KOTTI_DISABLE_POPULATORS"] = "1"
      conf_defaults["kotti.root_factory"] = [lambda req: None]
  
      def callback(arguments):
          args = ()
          args_with_location = (arguments["--scripts"] or DEFAULT_LOCATION,)
          if arguments["list_all"]:
              func = list_all
          elif arguments["upgrade"]:
              func = upgrade
              args = args_with_location + (arguments["--rev"],)
          elif arguments["upgrade_all"]:
              func = upgrade_all
          elif arguments["stamp_head"]:
              func = stamp_head
              args = args_with_location + (arguments["--rev"],)
          else:
              raise ValueError("Unknown command")
          func(*args)
  
      try:
          return command(callback, __doc__)
      finally:
          conf_defaults.clear()
          conf_defaults.update(save_conf_defaults)
          del os.environ["KOTTI_DISABLE_POPULATORS"]
  