<commit_before>import argparse
import os
import sys
import textwrap
from typing import Dict

import configargparse

import locust

version = locust.__version__


DEFAULT_CONFIG_FILES = ["~/.locust.conf", "locust.conf"]


class LocustArgumentParser(configargparse.ArgumentParser):
    """Drop-in replacement for `configargparse.ArgumentParser` that adds support for
    optionally exclude arguments from the UI.
    """

    def add_argument(self, *args, **kwargs) -> configargparse.Action:
        """
        This method supports the same args as ArgumentParser.add_argument(..)
        as well as the additional args below.

        Arguments:
            include_in_web_ui: If True (default), the argument will show in the UI.

        Returns:
            argparse.Action: the new argparse action
        """
        include_in_web_ui = kwargs.pop("include_in_web_ui", True)
        action = super().add_argument(*args, **kwargs)
        action.include_in_web_ui = include_in_web_ui
        return action

    @property
    def args_included_in_web_ui(self) -> Dict[str, configargparse.Action]:
        return {a.dest: a for a in self._actions if hasattr(a, "include_in_web_ui") and a.include_in_web_ui}


def _is_package(path):
    """
    Is the given path a Python package?
    """
    return os.path.isdir(path) and os.path.exists(os.path.join(path, "__init__.py"))


def find_locustfile(locustfile):
    """
    Attempt to locate a locustfile, either explicitly or by searching parent dirs.
    """
    # Obtain env value
    names = [locustfile]
    # Create .py version if necessary
    if not names[0].endswith(".py"):
        names.append(names[0] + ".py")
    # Does the name contain path elements?
    if os.path.dirname(names[0]):
        # If so, expand home-directory markers and test for existence
        for name in names:
            expanded = os.path.expanduser(name)
            if os.path.exists(expanded):
                if name.endswith(".py") or _is_package(expanded):
                    return os.path.abspath(expanded)
    else:
        # Otherwise, start in cwd and work downwards towards filesystem root
        path = os.path.abspath(".")
        while True:
            for name in names:
                joined = os.path.join(path, name)
                if os.path.exists(joined):
                    if name.endswith(".py") or _is_package(joined):
                        return os.path.abspath(joined)
            parent_path = os.path.dirname(path)
            if parent_path == path:
                # we've reached the root path which has been checked this iteration
                break
            path = parent_path
    # Implicit 'return None' if nothing was found


def get_empty_argument_parser(add_help=True, default_config_files=DEFAULT_CONFIG_FILES) -> LocustArgumentParser:
    parser = LocustArgumentParser(
        default_config_files=default_config_files,
        add_env_var_help=False,
        add_config_file_help=False,
        add_help=add_help,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        usage=argparse.SUPPRESS,
        description=textwrap.dedent(
            """
            Usage: locust [OPTIONS] [UserClass ...]

        """
        ),
        # epilog="",
    )
    parser.add_argument(
        "-f",
        "--locustfile",
        default="locustfile",
        help="Python module file to import, e.g. '../other.py'. Default: locustfile",
        env_var="LOCUST_LOCUSTFILE",
    )
    parser.add_argument("--config", is_config_file_arg=True, help="Config file path")

    return parser


def parse_locustfile_option(args=None):
    """
    Construct a command line parser that is only used to parse the -f argument so that we can
    import the test scripts in case any of them adds additional command line arguments to the
    parser
    """
    parser = get_empty_argument_parser(add_help=False)
    parser.add_argument(
        "-h",
        "--help",
        action="store_true",
        default=False,
    )
    parser.add_argument(
        "--version",
        "-V",
        action="store_true",
        default=False,
    )

    options, _ = parser.parse_known_args(args=args)

    locustfile = find_locustfile(options.locustfile)

    if not locustfile:
        if options.help or options.version:
            # if --help or --version is specified we'll call parse_options which will print the help/version message
            parse_options(args=args)
        sys.stderr.write(
            "Could not find any locustfile! Ensure file ends in '.py' and see --help for available options.\n"
        )
        sys.exit(1)

    if locustfile == "locust.py":<commit_msg>Improve error message if no locust file can be found.<commit_after>import argparse
import os
import sys
import textwrap
from typing import Dict

import configargparse

import locust

version = locust.__version__


DEFAULT_CONFIG_FILES = ["~/.locust.conf", "locust.conf"]


class LocustArgumentParser(configargparse.ArgumentParser):
    """Drop-in replacement for `configargparse.ArgumentParser` that adds support for
    optionally exclude arguments from the UI.
    """

    def add_argument(self, *args, **kwargs) -> configargparse.Action:
        """
        This method supports the same args as ArgumentParser.add_argument(..)
        as well as the additional args below.

        Arguments:
            include_in_web_ui: If True (default), the argument will show in the UI.

        Returns:
            argparse.Action: the new argparse action
        """
        include_in_web_ui = kwargs.pop("include_in_web_ui", True)
        action = super().add_argument(*args, **kwargs)
        action.include_in_web_ui = include_in_web_ui
        return action

    @property
    def args_included_in_web_ui(self) -> Dict[str, configargparse.Action]:
        return {a.dest: a for a in self._actions if hasattr(a, "include_in_web_ui") and a.include_in_web_ui}


def _is_package(path):
    """
    Is the given path a Python package?
    """
    return os.path.isdir(path) and os.path.exists(os.path.join(path, "__init__.py"))


def find_locustfile(locustfile):
    """
    Attempt to locate a locustfile, either explicitly or by searching parent dirs.
    """
    # Obtain env value
    names = [locustfile]
    # Create .py version if necessary
    if not names[0].endswith(".py"):
        names.append(names[0] + ".py")
    # Does the name contain path elements?
    if os.path.dirname(names[0]):
        # If so, expand home-directory markers and test for existence
        for name in names:
            expanded = os.path.expanduser(name)
            if os.path.exists(expanded):
                if name.endswith(".py") or _is_package(expanded):
                    return os.path.abspath(expanded)
    else:
        # Otherwise, start in cwd and work downwards towards filesystem root
        path = os.path.abspath(".")
        while True:
            for name in names:
                joined = os.path.join(path, name)
                if os.path.exists(joined):
                    if name.endswith(".py") or _is_package(joined):
                        return os.path.abspath(joined)
            parent_path = os.path.dirname(path)
            if parent_path == path:
                # we've reached the root path which has been checked this iteration
                break
            path = parent_path
    # Implicit 'return None' if nothing was found


def get_empty_argument_parser(add_help=True, default_config_files=DEFAULT_CONFIG_FILES) -> LocustArgumentParser:
    parser = LocustArgumentParser(
        default_config_files=default_config_files,
        add_env_var_help=False,
        add_config_file_help=False,
        add_help=add_help,
        formatter_class=argparse.RawDescriptionHelpFormatter,
        usage=argparse.SUPPRESS,
        description=textwrap.dedent(
            """
            Usage: locust [OPTIONS] [UserClass ...]

        """
        ),
        # epilog="",
    )
    parser.add_argument(
        "-f",
        "--locustfile",
        default="locustfile",
        help="Python module to import, e.g. '../other_test.py'. Either a .py file or a package directory. Defaults to 'locustfile'",
        env_var="LOCUST_LOCUSTFILE",
    )
    parser.add_argument("--config", is_config_file_arg=True, help="Config file path")

    return parser


def parse_locustfile_option(args=None):
    """
    Construct a command line parser that is only used to parse the -f argument so that we can
    import the test scripts in case any of them adds additional command line arguments to the
    parser
    """
    parser = get_empty_argument_parser(add_help=False)
    parser.add_argument(
        "-h",
        "--help",
        action="store_true",
        default=False,
    )
    parser.add_argument(
        "--version",
        "-V",
        action="store_true",
        default=False,
    )

    options, _ = parser.parse_known_args(args=args)

    locustfile = find_locustfile(options.locustfile)

    if not locustfile:
        if options.help or options.version:
            # if --help or --version is specified we'll call parse_options which will print the help/version message
            parse_options(args=args)
        if options.locustfile.endswith(".py") or options.locustfile == "locustfile":
            note_about_file_endings = ""
        else:
            note_about_file_endings = "Ensure your file ends with '.py'. "
        sys.stderr.write(
            f"Could not find any locustfile called '{options.locustfile}'. {note_about_file_endings}See --help for available options.\n"
        )
        sys.exit(1)

    if locustfile == "locust.py":
{'fraction_non_alphanumeric': 0.09590739975196362, 'fraction_numerical': 0.0008267879288962381, 'mean_word_length': 2.435214767483138, 'pattern_counts': {'":': 3, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1328031', 'n_tokens_mistral': 2838, 'n_tokens_neox': 2772, 'n_words': 762}