<commit_before>//! Complete commands within shells

/// Complete commands within bash
pub mod bash {
    use std::ffi::OsString;
    use std::io::Write;

    use unicode_xid::UnicodeXID;

    #[derive(clap::Subcommand)]
    #[clap(hide = true)]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub enum CompleteCommand {
        /// Register shell completions for this program
        Complete(CompleteArgs),
    }

    #[derive(clap::Args)]
    #[clap(group = clap::ArgGroup::new("complete").multiple(true).conflicts_with("register"))]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub struct CompleteArgs {
        /// Path to write completion-registration to
        #[clap(long, required = true, parse(from_os_str))]
        register: Option<std::path::PathBuf>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        index: Option<usize>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        ifs: Option<String>,

        #[clap(
            long = "type",
            required = true,
            arg_enum,
            hide_short_help = true,
            group = "complete"
        )]
        comp_type: Option<CompType>,

        #[clap(long, hide_short_help = true, group = "complete")]
        space: bool,

        #[clap(
            long,
            conflicts_with = "space",
            hide_short_help = true,
            group = "complete"
        )]
        no_space: bool,

        #[clap(raw = true, hide_short_help = true, group = "complete")]
        comp_words: Vec<OsString>,
    }

    impl CompleteCommand {
        /// Process the completion request
        pub fn complete(&self, cmd: &mut clap::Command) -> std::convert::Infallible {
            self.try_complete(cmd).unwrap_or_else(|e| e.exit());
            std::process::exit(0)
        }

        /// Process the completion request
        pub fn try_complete(&self, cmd: &mut clap::Command) -> clap::Result<()> {
            debug!("CompleteCommand::try_complete: {:?}", self);
            let CompleteCommand::Complete(args) = self;
            if let Some(out_path) = args.register.as_deref() {
                let mut buf = Vec::new();
                let name = cmd.get_name();
                let bin = cmd.get_bin_name().unwrap_or(cmd.get_name());
                register(name, [bin], bin, &Behavior::default(), &mut buf)?;
                if out_path == std::path::Path::new("-") {
                    std::io::stdout().write(&buf)?;
                } else {
                    if out_path.is_dir() {
                        let out_path = out_path.join(file_name(name));
                        std::fs::write(out_path, buf)?;
                    } else {
                        std::fs::write(out_path, buf)?;
                    }
                }
            } else {
                let index = args.index.unwrap_or_default();
                let comp_type = args.comp_type.unwrap_or_default();
                let space = match (args.space, args.no_space) {
                    (true, false) => Some(true),
                    (false, true) => Some(false),
                    (true, true) => {
                        unreachable!("`--space` and `--no-space` set, clap should prevent this")
                    }
                    (false, false) => None,
                }
                .unwrap();
                let current_dir = std::env::current_dir().ok();
                let completions = complete(
                    cmd,
                    args.comp_words.clone(),
                    index,
                    comp_type,
                    space,
                    current_dir.as_deref(),
                )?;

                let mut buf = Vec::new();
                for completion in &completions {
                    write!(&mut buf, "{}", completion.to_string_lossy())?;
                    write!(&mut buf, "{}", args.ifs.as_deref().unwrap_or("\n"))?;
                }
                std::io::stdout().write(&buf)?;
            }

            Ok(())
        }
    }

    /// The recommended file name for the registration code
    pub fn file_name(name: &str) -> String {
        format!("{}.bash", name)
    }

    /// Define the completion behavior
    pub enum Behavior {
        /// Bare bones behavior
        Minimal,
        /// Fallback to readline behavior when no matches are generated
        Readline,
        /// Customize bash's completion behavior
        Custom(String),
    }

    impl Default for Behavior {
        fn default() -> Self {
            Self::Readline
        }
    }

    /// Generate code to register the dynamic completion
    pub fn register(
        name: &str,
        executables: impl IntoIterator<Item = impl AsRef<str>>,
        completer: &str,
        behavior: &Behavior,
        buf: &mut dyn Write,
    ) -> Result<(), std::io::Error> {
        let escaped_name = name.replace("-", "_");
        debug_assert!(
            escaped_name.chars().all(|c| c.is_xid_continue()),
            "`name` must be an identifier, got `{}`",
            escaped_name
        );
        let mut upper_name = escaped_name.clone();
        upper_name.make_ascii_uppercase();

        let executables = executables
            .into_iter()
            .map(|s| shlex::quote(s.as_ref()).into_owned())
            .collect::<Vec<_>>()
            .join(" ");

        let options = match behavior {
            Behavior::Minimal => "-o nospace -o bashdefault",
            Behavior::Readline => "-o nospace -o default -o bashdefault",
            Behavior::Custom(c) => c.as_str(),
        };

        let completer = shlex::quote(completer);

        let script = r#"
_clap_complete_NAME() {
    local IFS=$'\013'
    local SUPPRESS_SPACE=0
    if compopt +o nospace 2> /dev/null; then
        SUPPRESS_SPACE=1
    fi
    if [[ ${SUPPRESS_SPACE} == 1 ]]; then
        SPACE_ARG="--no-space"
    else
        SPACE_ARG="--space"
    fi
    COMPREPLY=( $("COMPLETER" complete --index ${COMP_CWORD} --type ${COMP_TYPE} ${SPACE_ARG} --ifs="$IFS" -- "${COMP_WORDS[@]}") )
    if [[ $? != 0 ]]; then
        unset COMPREPLY
    elif [[ $SUPPRESS_SPACE == 1 ]] && [[ "${COMPREPLY-}" =~ [=/:]$ ]]; then
        compopt -o nospace
    fi
}
complete OPTIONS -F _clap_complete_NAME EXECUTABLES
"#
        .replace("NAME", &escaped_name)
        .replace("EXECUTABLES", &executables)
        .replace("OPTIONS", options)
        .replace("COMPLETER", &completer)
        .replace("UPPER", &upper_name);

        writeln!(buf, "{}", script)?;
        Ok(())
    }

    /// Type of completion attempted that caused a completion function to be called
    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
    #[non_exhaustive]
    pub enum CompType {
        /// Normal completion
        Normal,
        /// List completions after successive tabs
        Successive,
        /// List alternatives on partial word completion
        Alternatives,
        /// List completions if the word is not unmodified
        Unmodified,
        /// Menu completion
        Menu,
    }

    impl clap::ArgEnum for CompType {
        fn value_variants<'a>() -> &'a [Self] {
            &[
                Self::Normal,
                Self::Successive,
                Self::Alternatives,
                Self::Unmodified,
                Self::Menu,
            ]
        }
        fn to_possible_value<'a>(&self) -> ::std::option::Option<clap::PossibleValue<'a>> {
            match self {
                Self::Normal => {
                    let value = "9";
                    debug_assert_eq!(b'\t'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("normal")
                            .help("Normal completion"),
                    )
                }
                Self::Successive => {
                    let value = "63";
                    debug_assert_eq!(b'?'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("successive")
                            .help("List completions after successive tabs"),
                    )
                }
                Self::Alternatives => {
                    let value = "33";
                    debug_assert_eq!(b'!'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("alternatives")
                            .help("List alternatives on partial word completion"),
                    )
                }
                Self::Unmodified => {
                    let value = "64";
                    debug_assert_eq!(b'@'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("unmodified")
                            .help("List completions if the word is not unmodified"),
                    )
                }
                Self::Menu => {
                    let value = "37";
                    debug_assert_eq!(b'%'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("menu")
                            .help("Menu completion"),
                    )
                }
            }
        }
    }

    impl Default for CompType {
        fn default() -> Self {
            Self::Normal
        }
    }

    /// Complete the command specified
    pub fn complete(
        cmd: &mut clap::Command,
        args: Vec<std::ffi::OsString>,
        arg_index: usize,
        _comp_type: CompType,
        _trailing_space: bool,
        current_dir: Option<&std::path::Path>,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        cmd.build();

        let raw_args = clap_lex::RawArgs::new(args.into_iter());
        let mut cursor = raw_args.cursor();
        let mut target_cursor = raw_args.cursor();
        raw_args.seek(
            &mut target_cursor,
            clap_lex::SeekFrom::Start(arg_index as u64),
        );
        // As we loop, `cursor` will always be pointing to the next item
        raw_args.next_os(&mut target_cursor);

        // TODO: Multicall support
        if !cmd.is_no_binary_name_set() {
            raw_args.next_os(&mut cursor);
        }

        let mut current_cmd = &*cmd;
        while let Some(arg) = raw_args.next(&mut cursor) {
            if cursor == target_cursor {
                return complete_new(current_cmd, current_dir);
            }
            if let Ok(value) = arg.to_value() {
                if let Some(next_cmd) = current_cmd.find_subcommand(value) {
                    current_cmd = next_cmd;
                }
            }
        }

        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "No completion generated",
        ))
    }

    fn complete_new(
        cmd: &clap::Command,
        current_dir: Option<&std::path::Path>,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        let mut completions = Vec::new();

        completions.extend(
            crate::generator::utils::longs_and_visible_aliases(cmd)
                .into_iter()
                .map(|f| format!("--{}", f).into()),
        );
        completions.extend(
            crate::generator::utils::shorts_and_visible_aliases(cmd)
                .into_iter()
                .map(|f| format!("-{}", f).into()),
        );

        let mut positionals = Vec::new();
        positionals.extend(
            cmd.get_positionals()
                .flat_map(|p| p.get_possible_values())
                .flatten()
                .map(|p| p.get_name().into()),
        );
        let hints = cmd
            .get_positionals()
            .map(|p| p.get_value_hint())
            .map(|h| {
                if h == clap::ValueHint::Unknown {
                    clap::ValueHint::AnyPath
                } else {
                    h
                }
            })
            .collect::<std::collections::HashSet<_>>();
        for hint in hints {
            match hint {
                clap::ValueHint::Unknown => unreachable!("Filtered out"),
                clap::ValueHint::Other => {
                    // Should not complete
                }
                clap::ValueHint::AnyPath => {
                    positionals.extend(complete_path(current_dir, |_| true));
                }
                clap::ValueHint::FilePath => {
                    positionals.extend(complete_path(current_dir, |p| p.is_file()));
                }
                clap::ValueHint::DirPath => {
                    positionals.extend(complete_path(current_dir, |p| p.is_dir()));
                }
                clap::ValueHint::ExecutablePath => {
                    use is_executable::IsExecutable;
                    positionals.extend(complete_path(current_dir, |p| p.is_executable()));
                }
                clap::ValueHint::CommandName
                | clap::ValueHint::CommandString
                | clap::ValueHint::CommandWithArguments
                | clap::ValueHint::Username
                | clap::ValueHint::Hostname
                | clap::ValueHint::Url
                | clap::ValueHint::EmailAddress => {
                    // No completion implementation
                }
                _ => {
                    // Safe-ish fallback
                    positionals.extend(complete_path(current_dir, |_| true));
                }
            }
        }
        positionals.sort();
        positionals.dedup();
        completions.extend(positionals);

        completions.extend(all_subcommands(cmd).into_iter().map(|s| s));

        Ok(completions)
    }

    fn complete_path(
        current_dir: Option<&std::path::Path>,
        is_wanted: impl Fn(&std::path::Path) -> bool,
    ) -> Vec<OsString> {
        let mut completions = Vec::new();

        let current_dir = match current_dir {
            Some(current_dir) => current_dir,
            None => {
                return Vec::new();
            }
        };
        for entry in std::fs::read_dir(current_dir)
            .ok()
            .into_iter()
            .flatten()
            .filter_map(Result::ok)
        {
            let mut path = entry.path();
            if entry.metadata().map(|m| m.is_dir()).unwrap_or(false) {
                path.push(""); // Ensure trailing `/`
                let suggestion = path.strip_prefix(current_dir).unwrap();
                completions.push(suggestion.as_os_str().to_owned());
            } else {
                if is_wanted(&path) {
                    let suggestion = path.strip_prefix(current_dir).unwrap();
                    completions.push(suggestion.as_os_str().to_owned());
                }
            }
        }

        completions
    }

    fn all_subcommands(cmd: &clap::Command) -> Vec<OsString> {
        debug!("all_subcommands");

        let mut scs = crate::generator::utils::all_subcommands(cmd)
            .iter()
            .map(|x| OsString::from(&x.0))
            .collect::<Vec<_>>();
        scs.sort();
        scs.dedup();
        scs
    }
}
<commit_msg>feat(complete): Only list for current positional<commit_after>//! Complete commands within shells

/// Complete commands within bash
pub mod bash {
    use std::ffi::OsString;
    use std::io::Write;

    use unicode_xid::UnicodeXID;

    #[derive(clap::Subcommand)]
    #[clap(hide = true)]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub enum CompleteCommand {
        /// Register shell completions for this program
        Complete(CompleteArgs),
    }

    #[derive(clap::Args)]
    #[clap(group = clap::ArgGroup::new("complete").multiple(true).conflicts_with("register"))]
    #[allow(missing_docs)]
    #[derive(Clone, Debug)]
    pub struct CompleteArgs {
        /// Path to write completion-registration to
        #[clap(long, required = true, parse(from_os_str))]
        register: Option<std::path::PathBuf>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        index: Option<usize>,

        #[clap(
            long,
            required = true,
            value_name = "COMP_CWORD",
            hide_short_help = true,
            group = "complete"
        )]
        ifs: Option<String>,

        #[clap(
            long = "type",
            required = true,
            arg_enum,
            hide_short_help = true,
            group = "complete"
        )]
        comp_type: Option<CompType>,

        #[clap(long, hide_short_help = true, group = "complete")]
        space: bool,

        #[clap(
            long,
            conflicts_with = "space",
            hide_short_help = true,
            group = "complete"
        )]
        no_space: bool,

        #[clap(raw = true, hide_short_help = true, group = "complete")]
        comp_words: Vec<OsString>,
    }

    impl CompleteCommand {
        /// Process the completion request
        pub fn complete(&self, cmd: &mut clap::Command) -> std::convert::Infallible {
            self.try_complete(cmd).unwrap_or_else(|e| e.exit());
            std::process::exit(0)
        }

        /// Process the completion request
        pub fn try_complete(&self, cmd: &mut clap::Command) -> clap::Result<()> {
            debug!("CompleteCommand::try_complete: {:?}", self);
            let CompleteCommand::Complete(args) = self;
            if let Some(out_path) = args.register.as_deref() {
                let mut buf = Vec::new();
                let name = cmd.get_name();
                let bin = cmd.get_bin_name().unwrap_or(cmd.get_name());
                register(name, [bin], bin, &Behavior::default(), &mut buf)?;
                if out_path == std::path::Path::new("-") {
                    std::io::stdout().write(&buf)?;
                } else {
                    if out_path.is_dir() {
                        let out_path = out_path.join(file_name(name));
                        std::fs::write(out_path, buf)?;
                    } else {
                        std::fs::write(out_path, buf)?;
                    }
                }
            } else {
                let index = args.index.unwrap_or_default();
                let comp_type = args.comp_type.unwrap_or_default();
                let space = match (args.space, args.no_space) {
                    (true, false) => Some(true),
                    (false, true) => Some(false),
                    (true, true) => {
                        unreachable!("`--space` and `--no-space` set, clap should prevent this")
                    }
                    (false, false) => None,
                }
                .unwrap();
                let current_dir = std::env::current_dir().ok();
                let completions = complete(
                    cmd,
                    args.comp_words.clone(),
                    index,
                    comp_type,
                    space,
                    current_dir.as_deref(),
                )?;

                let mut buf = Vec::new();
                for completion in &completions {
                    write!(&mut buf, "{}", completion.to_string_lossy())?;
                    write!(&mut buf, "{}", args.ifs.as_deref().unwrap_or("\n"))?;
                }
                std::io::stdout().write(&buf)?;
            }

            Ok(())
        }
    }

    /// The recommended file name for the registration code
    pub fn file_name(name: &str) -> String {
        format!("{}.bash", name)
    }

    /// Define the completion behavior
    pub enum Behavior {
        /// Bare bones behavior
        Minimal,
        /// Fallback to readline behavior when no matches are generated
        Readline,
        /// Customize bash's completion behavior
        Custom(String),
    }

    impl Default for Behavior {
        fn default() -> Self {
            Self::Readline
        }
    }

    /// Generate code to register the dynamic completion
    pub fn register(
        name: &str,
        executables: impl IntoIterator<Item = impl AsRef<str>>,
        completer: &str,
        behavior: &Behavior,
        buf: &mut dyn Write,
    ) -> Result<(), std::io::Error> {
        let escaped_name = name.replace("-", "_");
        debug_assert!(
            escaped_name.chars().all(|c| c.is_xid_continue()),
            "`name` must be an identifier, got `{}`",
            escaped_name
        );
        let mut upper_name = escaped_name.clone();
        upper_name.make_ascii_uppercase();

        let executables = executables
            .into_iter()
            .map(|s| shlex::quote(s.as_ref()).into_owned())
            .collect::<Vec<_>>()
            .join(" ");

        let options = match behavior {
            Behavior::Minimal => "-o nospace -o bashdefault",
            Behavior::Readline => "-o nospace -o default -o bashdefault",
            Behavior::Custom(c) => c.as_str(),
        };

        let completer = shlex::quote(completer);

        let script = r#"
_clap_complete_NAME() {
    local IFS=$'\013'
    local SUPPRESS_SPACE=0
    if compopt +o nospace 2> /dev/null; then
        SUPPRESS_SPACE=1
    fi
    if [[ ${SUPPRESS_SPACE} == 1 ]]; then
        SPACE_ARG="--no-space"
    else
        SPACE_ARG="--space"
    fi
    COMPREPLY=( $("COMPLETER" complete --index ${COMP_CWORD} --type ${COMP_TYPE} ${SPACE_ARG} --ifs="$IFS" -- "${COMP_WORDS[@]}") )
    if [[ $? != 0 ]]; then
        unset COMPREPLY
    elif [[ $SUPPRESS_SPACE == 1 ]] && [[ "${COMPREPLY-}" =~ [=/:]$ ]]; then
        compopt -o nospace
    fi
}
complete OPTIONS -F _clap_complete_NAME EXECUTABLES
"#
        .replace("NAME", &escaped_name)
        .replace("EXECUTABLES", &executables)
        .replace("OPTIONS", options)
        .replace("COMPLETER", &completer)
        .replace("UPPER", &upper_name);

        writeln!(buf, "{}", script)?;
        Ok(())
    }

    /// Type of completion attempted that caused a completion function to be called
    #[derive(Copy, Clone, Debug, PartialEq, Eq)]
    #[non_exhaustive]
    pub enum CompType {
        /// Normal completion
        Normal,
        /// List completions after successive tabs
        Successive,
        /// List alternatives on partial word completion
        Alternatives,
        /// List completions if the word is not unmodified
        Unmodified,
        /// Menu completion
        Menu,
    }

    impl clap::ArgEnum for CompType {
        fn value_variants<'a>() -> &'a [Self] {
            &[
                Self::Normal,
                Self::Successive,
                Self::Alternatives,
                Self::Unmodified,
                Self::Menu,
            ]
        }
        fn to_possible_value<'a>(&self) -> ::std::option::Option<clap::PossibleValue<'a>> {
            match self {
                Self::Normal => {
                    let value = "9";
                    debug_assert_eq!(b'\t'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("normal")
                            .help("Normal completion"),
                    )
                }
                Self::Successive => {
                    let value = "63";
                    debug_assert_eq!(b'?'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("successive")
                            .help("List completions after successive tabs"),
                    )
                }
                Self::Alternatives => {
                    let value = "33";
                    debug_assert_eq!(b'!'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("alternatives")
                            .help("List alternatives on partial word completion"),
                    )
                }
                Self::Unmodified => {
                    let value = "64";
                    debug_assert_eq!(b'@'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("unmodified")
                            .help("List completions if the word is not unmodified"),
                    )
                }
                Self::Menu => {
                    let value = "37";
                    debug_assert_eq!(b'%'.to_string(), value);
                    Some(
                        clap::PossibleValue::new(value)
                            .alias("menu")
                            .help("Menu completion"),
                    )
                }
            }
        }
    }

    impl Default for CompType {
        fn default() -> Self {
            Self::Normal
        }
    }

    /// Complete the command specified
    pub fn complete(
        cmd: &mut clap::Command,
        args: Vec<std::ffi::OsString>,
        arg_index: usize,
        _comp_type: CompType,
        _trailing_space: bool,
        current_dir: Option<&std::path::Path>,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        cmd.build();

        let raw_args = clap_lex::RawArgs::new(args.into_iter());
        let mut cursor = raw_args.cursor();
        let mut target_cursor = raw_args.cursor();
        raw_args.seek(
            &mut target_cursor,
            clap_lex::SeekFrom::Start(arg_index as u64),
        );
        // As we loop, `cursor` will always be pointing to the next item
        raw_args.next_os(&mut target_cursor);

        // TODO: Multicall support
        if !cmd.is_no_binary_name_set() {
            raw_args.next_os(&mut cursor);
        }

        let mut current_cmd = &*cmd;
        let mut pos_index = 0;
        while let Some(arg) = raw_args.next(&mut cursor) {
            if cursor == target_cursor {
                return complete_new(current_cmd, current_dir, pos_index);
            }
            if let Ok(value) = arg.to_value() {
                if let Some(next_cmd) = current_cmd.find_subcommand(value) {
                    current_cmd = next_cmd;
                    pos_index = 0;
                    continue;
                }
            }

            pos_index += 1;
        }

        Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            "No completion generated",
        ))
    }

    fn complete_new(
        cmd: &clap::Command,
        current_dir: Option<&std::path::Path>,
        pos_index: usize,
    ) -> Result<Vec<std::ffi::OsString>, std::io::Error> {
        let mut completions = Vec::new();

        completions.extend(
            crate::generator::utils::longs_and_visible_aliases(cmd)
                .into_iter()
                .map(|f| format!("--{}", f).into()),
        );
        completions.extend(
            crate::generator::utils::shorts_and_visible_aliases(cmd)
                .into_iter()
                .map(|f| format!("-{}", f).into()),
        );

        let mut positionals = Vec::new();
        if let Some(positional) = cmd
            .get_positionals()
            .find(|p| p.get_index() == Some(pos_index))
        {
            if let Some(possible_values) = positional.get_possible_values() {
                positionals.extend(possible_values.into_iter().map(|p| p.get_name().into()));
            } else {
                match positional.get_value_hint() {
                    clap::ValueHint::Other => {
                        // Should not complete
                    }
                    clap::ValueHint::Unknown | clap::ValueHint::AnyPath => {
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                    clap::ValueHint::FilePath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_file()));
                    }
                    clap::ValueHint::DirPath => {
                        positionals.extend(complete_path(current_dir, |p| p.is_dir()));
                    }
                    clap::ValueHint::ExecutablePath => {
                        use is_executable::IsExecutable;
                        positionals.extend(complete_path(current_dir, |p| p.is_executable()));
                    }
                    clap::ValueHint::CommandName
                    | clap::ValueHint::CommandString
                    | clap::ValueHint::CommandWithArguments
                    | clap::ValueHint::Username
                    | clap::ValueHint::Hostname
                    | clap::ValueHint::Url
                    | clap::ValueHint::EmailAddress => {
                        // No completion implementation
                    }
                    _ => {
                        // Safe-ish fallback
                        positionals.extend(complete_path(current_dir, |_| true));
                    }
                }
                positionals.sort();
            }
            completions.extend(positionals);
        }

        completions.extend(all_subcommands(cmd).into_iter().map(|s| s));

        Ok(completions)
    }

    fn complete_path(
        current_dir: Option<&std::path::Path>,
        is_wanted: impl Fn(&std::path::Path) -> bool,
    ) -> Vec<OsString> {
        let mut completions = Vec::new();

        let current_dir = match current_dir {
            Some(current_dir) => current_dir,
            None => {
                return Vec::new();
            }
        };
        for entry in std::fs::read_dir(current_dir)
            .ok()
            .into_iter()
            .flatten()
            .filter_map(Result::ok)
        {
            let mut path = entry.path();
            if entry.metadata().map(|m| m.is_dir()).unwrap_or(false) {
                path.push(""); // Ensure trailing `/`
                let suggestion = path.strip_prefix(current_dir).unwrap();
                completions.push(suggestion.as_os_str().to_owned());
            } else {
                if is_wanted(&path) {
                    let suggestion = path.strip_prefix(current_dir).unwrap();
                    completions.push(suggestion.as_os_str().to_owned());
                }
            }
        }

        completions
    }

    fn all_subcommands(cmd: &clap::Command) -> Vec<OsString> {
        debug!("all_subcommands");

        let mut scs = crate::generator::utils::all_subcommands(cmd)
            .iter()
            .map(|x| OsString::from(&x.0))
            .collect::<Vec<_>>();
        scs.sort();
        scs.dedup();
        scs
    }
}
<|endoftext|>


<commit_before>use std::collections::BTreeSet;
use std::collections::{HashMap};
use std::io::{self, stdout, stderr, Write};
use std::path::Path;

use time::{SteadyTime, Duration};
use argparse::{ArgumentParser, List};
use signal::trap::Trap;
use nix::sys::signal::{SIGINT, SIGTERM, SIGCHLD, SIGTTIN, SIGTTOU, SIGCONT};
use nix::sys::signal::{SIGQUIT, SIGKILL};
use unshare::{Command, Namespace, reap_zombies};

use options::build_mode::{build_mode, BuildMode};
use container::mount::{mount_tmpfs};
use container::nsutil::{set_namespace, unshare_namespace};
use config::{Settings};
use config::command::{SuperviseInfo, Networking};
use config::command::SuperviseMode::{stop_on_failure};
use config::command::ChildCommand::{BridgeCommand};
use tty_util::{TtyGuard};
use super::network;
use super::build::{build_container};
use file_util::create_dir;
use process_util::{convert_status, killpg};
use super::wrap::Wrapper;

pub fn run_supervise_command(settings: &Settings, workdir: &Path,
    sup: &SuperviseInfo, cmdname: String, mut args: Vec<String>,
    mut bmode: BuildMode)
    -> Result<i32, String>
{
    let mut only: Vec<String> = Vec::new();
    let mut exclude: Vec<String> = Vec::new();
    if sup.mode != stop_on_failure {
        panic!("Only stop-on-failure mode implemented");
    }
    {
        args.insert(0, "vagga ".to_string() + &cmdname);
        let mut ap = ArgumentParser::new();
        ap.set_description(sup.description.as_ref().map(|x| &x[..])
            .unwrap_or("Run multiple processes simultaneously"));
        ap.refer(&mut only).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--only"], List, "
                Only run specified processes.
                This matches both names and tags");
        ap.refer(&mut exclude).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--exclude"], List, "
                Don't run specified processes.
                This excludes both names and tags");
        build_mode(&mut ap, &mut bmode);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }

    let mut containers = BTreeSet::new();
    let mut containers_in_netns = vec!();
    let mut bridges = vec!();
    let mut containers_host_net = vec!();
    let mut forwards = vec!();
    let mut ports = vec!();
    let mut versions = HashMap::new();
    let filtered_children = sup.children
        .iter().filter(|&(ref name, ref child)| {
            if only.len() > 0 {
                only.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_some()
            } else {
                exclude.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_none()
            }
        });
    for (name, child) in filtered_children {
        let cont = child.get_container();
        if !containers.contains(cont) {
            containers.insert(cont.to_string());
            let ver = try!(build_container(settings, cont, bmode));
            versions.insert(cont.to_string(), ver);
        }
        if let &BridgeCommand(_) = child {
            bridges.push(name.to_string());
        } else {
            if let Some(ref netw) = child.network() {
                containers_in_netns.push(name.to_string());
                for (ext_port, int_port) in netw.ports.iter() {
                     forwards.push((*ext_port, netw.ip.clone(), *int_port));
                    ports.push(*ext_port);
                }
            } else {
                containers_host_net.push(name.to_string());
            }
        }
    }
    containers_in_netns.extend(bridges.into_iter()); // Bridges are just last
    if containers_in_netns.len() > 0 && !network::is_netns_set_up() {
        return Err(format!("Network namespace is not set up. You need to run \
            vagga _create_netns first"));
    }
    debug!("Containers {} with host neworking, {} in netns",
        containers_host_net.len(), containers_in_netns.len());

    // Trap must be installed before tty_guard because TTY guard relies on
    // SIGTTOU and SIGTTIN be masked out
    let mut trap = Trap::trap(&[SIGINT, SIGQUIT,
                                SIGTERM, SIGCHLD, SIGTTOU, SIGTTIN]);
    let mut tty_guard = try!(TtyGuard::capture_tty()
        .map_err(|e| format!("Error handling tty: {}", e)));

    let mut children = HashMap::new();
    let mut error = false;
    for name in containers_host_net.iter() {
        let mut cmd: Command = Wrapper::new(
            Some(&versions[sup.children[name].get_container()]),
            settings);
        cmd.workdir(workdir);
        cmd.userns();
        cmd.arg(&cmdname);
        cmd.arg(&name);
        cmd.make_group_leader(true);
        match cmd.spawn() {
            Ok(child) => { children.insert(child.pid(), (name, child)); }
            Err(e) => {
                if !error {
                    println!(
                        "---------- \
                        Process {} could not be run: {}. Shutting down \
                        -----------",
                        name, e);
                    error = true;
                }
            }
        }
    }
    let port_forward_guard;
    if containers_in_netns.len() > 0 {
        let gwdir = network::namespace_dir();
        let nsdir = gwdir.join("children");
        if !nsdir.exists() {
            try_msg!(create_dir(&nsdir, false),
                     "Failed to create dir: {err}");
        }
        try!(network::join_gateway_namespaces());
        try!(unshare_namespace(Namespace::Mount)
            .map_err(|e| format!("Failed to create mount namespace: {}", e)));
        try!(mount_tmpfs(&nsdir, "size=10m,mode=0755"));

        let bridge_ns = nsdir.join("bridge");
        let ip = try!(network::setup_bridge(&bridge_ns, &forwards));

        port_forward_guard = network::PortForwardGuard::new(
            &gwdir.join("netns"), ip, ports);
        try!(port_forward_guard.start_forwarding());

        for name in containers_in_netns.iter() {
            let child = sup.children.get(name).unwrap();
            let mut cmd: Command = Wrapper::new(
                Some(&versions[sup.children[name].get_container()]),
                settings);
            cmd.workdir(workdir);
            cmd.arg(&cmdname);
            cmd.arg(&name);

            try!(set_namespace(&bridge_ns, Namespace::Net)
                .map_err(|e| format!("Error setting netns: {}", e)));
            if let &BridgeCommand(_) = child {
                // Already setup by set_namespace
                // But also need to mount namespace_dir into container
                cmd.env("VAGGA_NAMESPACE_DIR", &nsdir);
            } else {
                let netw = child.network().unwrap();
                let net_ns;
                let uts_ns;
                net_ns = nsdir.join("net.".to_string() + &netw.ip);
                uts_ns = nsdir.join("uts.".to_string() + &netw.ip);
                // TODO(tailhook) support multiple commands with same IP
                try!(network::setup_container(&net_ns, &uts_ns,
                    &name, &netw.ip,
                    &netw.hostname.as_ref().unwrap_or(name)));
                try!(set_namespace(&net_ns, Namespace::Net)
                    .map_err(|e| format!("Error setting netns: {}", e)));
                try!(set_namespace(&uts_ns, Namespace::Uts)
                    .map_err(|e| format!("Error setting netns: {}", e)));
            }

            cmd.make_group_leader(true);
            match cmd.spawn() {
                Ok(child) => { children.insert(child.pid(), (name, child)); }
                Err(e) => {
                    if !error {
                        println!(
                            "---------- \
                            Process {} could not be run: {}. Shutting down \
                            -----------",
                            name, e);
                        error = true;
                    }
                }
            }
        }

        // Need to set network namespace back to bridge, to keep namespace
        // alive. Otherwise bridge is dropped, and no connectivity between
        // containers.
        try!(set_namespace(&bridge_ns, Namespace::Net)
            .map_err(|e| format!("Error setting netns: {}", e)));
    }

    let mut errcode = 0;
    if error {
        errcode = 127;
        for &(_, ref child) in children.values() {
            child.signal(SIGTERM).ok();
        }
    } else {
        // Normal loop
        assert!(children.len() > 0);
        'signal_loop: for signal in trap.by_ref() {
            match signal {
                SIGINT => {
                    // SIGINT is usually a Ctrl+C, if we trap it here
                    // child process hasn't controlling terminal,
                    // so we send the signal to the child process
                    writeln!(&mut stderr(), "Received SIGINT signal. \
                        Waiting the processes to stop...").ok();
                    for &(cmd, ref child) in children.values() {
                        if unsafe { killpg(child.pid(), SIGTERM) } < 0 {
                             error!("Error sending SIGTERM to {:?}: {}", cmd,
                                io::Error::last_os_error());
                        }
                    }
                    errcode = 128+SIGINT;
                    break;
                }
                SIGTERM|SIGQUIT => {
                    // SIGTERM is usually sent to a specific process so we
                    // forward it to children
                    writeln!(&mut stderr(), "Received {} signal, \
                        propagating", signal).ok();
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                    errcode = 128+signal;
                    break;
                }
                SIGCHLD => {
                    for (pid, status) in reap_zombies() {
                        if let Some((name, _)) = children.remove(&pid) {
                            errcode = convert_status(status);
                            println!(
                                "---------- \
                                Process {}:{} {}. Shutting down \
                                -----------",
                                name, pid, status);
                            for (pid, _) in reap_zombies() {
                                children.remove(&pid);
                            }
                            for &(_, ref child) in children.values() {
                                child.signal(SIGTERM).ok();
                            }
                            break 'signal_loop;
                        }
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                _ => unreachable!(),
            }
        }
    }

    // Stopping loop
    if children.len() > 0 {
        let timeo = sup.kill_unresponsive_after;
        let mut deadline = SteadyTime::now() + Duration::seconds(timeo as i64);
        loop {
            match trap.wait(deadline) {
                Some(SIGINT) => {}
                Some(SIGTERM) => {
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                }
                Some(SIGCHLD) => {
                    for (pid, _) in reap_zombies() {
                        children.remove(&pid);
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                Some(_) => unreachable!(),
                None => {
                    println!(
                        "---------- \
                        Processes {:?} are still alive. Killing ... \
                        -----------\n\
                        To prevent killing of the processes set \
                        ``kill-unresponsive-after`` to some larger value",
                        children.values().map(|&(name, _)| name)
                            .collect::<Vec<_>>());
                    for &(_, ref child) in children.values() {
                        child.signal(SIGKILL).ok();
                    }
                    // Basically this deadline should never happen
                    deadline = SteadyTime::now() + Duration::seconds(3600);
                }
            }
        }
    }

    if errcode == 0 {
        if let Some(ref epilog) = sup.epilog {
            print!("{}", epilog);
        }
    }

    Ok(errcode)
}
<commit_msg>Fix unused warning<commit_after>use std::collections::BTreeSet;
use std::collections::{HashMap};
use std::io::{self, stdout, stderr, Write};
use std::path::Path;

use time::{SteadyTime, Duration};
use argparse::{ArgumentParser, List};
use signal::trap::Trap;
use nix::sys::signal::{SIGINT, SIGTERM, SIGCHLD, SIGTTIN, SIGTTOU};
use nix::sys::signal::{SIGQUIT, SIGKILL};
use unshare::{Command, Namespace, reap_zombies};

use options::build_mode::{build_mode, BuildMode};
use container::mount::{mount_tmpfs};
use container::nsutil::{set_namespace, unshare_namespace};
use config::{Settings};
use config::command::{SuperviseInfo, Networking};
use config::command::SuperviseMode::{stop_on_failure};
use config::command::ChildCommand::{BridgeCommand};
use tty_util::{TtyGuard};
use super::network;
use super::build::{build_container};
use file_util::create_dir;
use process_util::{convert_status, killpg};
use super::wrap::Wrapper;

pub fn run_supervise_command(settings: &Settings, workdir: &Path,
    sup: &SuperviseInfo, cmdname: String, mut args: Vec<String>,
    mut bmode: BuildMode)
    -> Result<i32, String>
{
    let mut only: Vec<String> = Vec::new();
    let mut exclude: Vec<String> = Vec::new();
    if sup.mode != stop_on_failure {
        panic!("Only stop-on-failure mode implemented");
    }
    {
        args.insert(0, "vagga ".to_string() + &cmdname);
        let mut ap = ArgumentParser::new();
        ap.set_description(sup.description.as_ref().map(|x| &x[..])
            .unwrap_or("Run multiple processes simultaneously"));
        ap.refer(&mut only).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--only"], List, "
                Only run specified processes.
                This matches both names and tags");
        ap.refer(&mut exclude).metavar("PROCESS_NAME_OR_TAG")
            .add_option(&["--exclude"], List, "
                Don't run specified processes.
                This excludes both names and tags");
        build_mode(&mut ap, &mut bmode);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }

    let mut containers = BTreeSet::new();
    let mut containers_in_netns = vec!();
    let mut bridges = vec!();
    let mut containers_host_net = vec!();
    let mut forwards = vec!();
    let mut ports = vec!();
    let mut versions = HashMap::new();
    let filtered_children = sup.children
        .iter().filter(|&(ref name, ref child)| {
            if only.len() > 0 {
                only.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_some()
            } else {
                exclude.iter().find(|x| {
                    name == x || child.get_tags().iter().any(|t| t == *x)
                }).is_none()
            }
        });
    for (name, child) in filtered_children {
        let cont = child.get_container();
        if !containers.contains(cont) {
            containers.insert(cont.to_string());
            let ver = try!(build_container(settings, cont, bmode));
            versions.insert(cont.to_string(), ver);
        }
        if let &BridgeCommand(_) = child {
            bridges.push(name.to_string());
        } else {
            if let Some(ref netw) = child.network() {
                containers_in_netns.push(name.to_string());
                for (ext_port, int_port) in netw.ports.iter() {
                     forwards.push((*ext_port, netw.ip.clone(), *int_port));
                    ports.push(*ext_port);
                }
            } else {
                containers_host_net.push(name.to_string());
            }
        }
    }
    containers_in_netns.extend(bridges.into_iter()); // Bridges are just last
    if containers_in_netns.len() > 0 && !network::is_netns_set_up() {
        return Err(format!("Network namespace is not set up. You need to run \
            vagga _create_netns first"));
    }
    debug!("Containers {} with host neworking, {} in netns",
        containers_host_net.len(), containers_in_netns.len());

    // Trap must be installed before tty_guard because TTY guard relies on
    // SIGTTOU and SIGTTIN be masked out
    let mut trap = Trap::trap(&[SIGINT, SIGQUIT,
                                SIGTERM, SIGCHLD, SIGTTOU, SIGTTIN]);
    let mut tty_guard = try!(TtyGuard::capture_tty()
        .map_err(|e| format!("Error handling tty: {}", e)));

    let mut children = HashMap::new();
    let mut error = false;
    for name in containers_host_net.iter() {
        let mut cmd: Command = Wrapper::new(
            Some(&versions[sup.children[name].get_container()]),
            settings);
        cmd.workdir(workdir);
        cmd.userns();
        cmd.arg(&cmdname);
        cmd.arg(&name);
        cmd.make_group_leader(true);
        match cmd.spawn() {
            Ok(child) => { children.insert(child.pid(), (name, child)); }
            Err(e) => {
                if !error {
                    println!(
                        "---------- \
                        Process {} could not be run: {}. Shutting down \
                        -----------",
                        name, e);
                    error = true;
                }
            }
        }
    }
    let port_forward_guard;
    if containers_in_netns.len() > 0 {
        let gwdir = network::namespace_dir();
        let nsdir = gwdir.join("children");
        if !nsdir.exists() {
            try_msg!(create_dir(&nsdir, false),
                     "Failed to create dir: {err}");
        }
        try!(network::join_gateway_namespaces());
        try!(unshare_namespace(Namespace::Mount)
            .map_err(|e| format!("Failed to create mount namespace: {}", e)));
        try!(mount_tmpfs(&nsdir, "size=10m,mode=0755"));

        let bridge_ns = nsdir.join("bridge");
        let ip = try!(network::setup_bridge(&bridge_ns, &forwards));

        port_forward_guard = network::PortForwardGuard::new(
            &gwdir.join("netns"), ip, ports);
        try!(port_forward_guard.start_forwarding());

        for name in containers_in_netns.iter() {
            let child = sup.children.get(name).unwrap();
            let mut cmd: Command = Wrapper::new(
                Some(&versions[sup.children[name].get_container()]),
                settings);
            cmd.workdir(workdir);
            cmd.arg(&cmdname);
            cmd.arg(&name);

            try!(set_namespace(&bridge_ns, Namespace::Net)
                .map_err(|e| format!("Error setting netns: {}", e)));
            if let &BridgeCommand(_) = child {
                // Already setup by set_namespace
                // But also need to mount namespace_dir into container
                cmd.env("VAGGA_NAMESPACE_DIR", &nsdir);
            } else {
                let netw = child.network().unwrap();
                let net_ns;
                let uts_ns;
                net_ns = nsdir.join("net.".to_string() + &netw.ip);
                uts_ns = nsdir.join("uts.".to_string() + &netw.ip);
                // TODO(tailhook) support multiple commands with same IP
                try!(network::setup_container(&net_ns, &uts_ns,
                    &name, &netw.ip,
                    &netw.hostname.as_ref().unwrap_or(name)));
                try!(set_namespace(&net_ns, Namespace::Net)
                    .map_err(|e| format!("Error setting netns: {}", e)));
                try!(set_namespace(&uts_ns, Namespace::Uts)
                    .map_err(|e| format!("Error setting netns: {}", e)));
            }

            cmd.make_group_leader(true);
            match cmd.spawn() {
                Ok(child) => { children.insert(child.pid(), (name, child)); }
                Err(e) => {
                    if !error {
                        println!(
                            "---------- \
                            Process {} could not be run: {}. Shutting down \
                            -----------",
                            name, e);
                        error = true;
                    }
                }
            }
        }

        // Need to set network namespace back to bridge, to keep namespace
        // alive. Otherwise bridge is dropped, and no connectivity between
        // containers.
        try!(set_namespace(&bridge_ns, Namespace::Net)
            .map_err(|e| format!("Error setting netns: {}", e)));
    }

    let mut errcode = 0;
    if error {
        errcode = 127;
        for &(_, ref child) in children.values() {
            child.signal(SIGTERM).ok();
        }
    } else {
        // Normal loop
        assert!(children.len() > 0);
        'signal_loop: for signal in trap.by_ref() {
            match signal {
                SIGINT => {
                    // SIGINT is usually a Ctrl+C, if we trap it here
                    // child process hasn't controlling terminal,
                    // so we send the signal to the child process
                    writeln!(&mut stderr(), "Received SIGINT signal. \
                        Waiting the processes to stop...").ok();
                    for &(cmd, ref child) in children.values() {
                        if unsafe { killpg(child.pid(), SIGTERM) } < 0 {
                             error!("Error sending SIGTERM to {:?}: {}", cmd,
                                io::Error::last_os_error());
                        }
                    }
                    errcode = 128+SIGINT;
                    break;
                }
                SIGTERM|SIGQUIT => {
                    // SIGTERM is usually sent to a specific process so we
                    // forward it to children
                    writeln!(&mut stderr(), "Received {} signal, \
                        propagating", signal).ok();
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                    errcode = 128+signal;
                    break;
                }
                SIGCHLD => {
                    for (pid, status) in reap_zombies() {
                        if let Some((name, _)) = children.remove(&pid) {
                            errcode = convert_status(status);
                            println!(
                                "---------- \
                                Process {}:{} {}. Shutting down \
                                -----------",
                                name, pid, status);
                            for (pid, _) in reap_zombies() {
                                children.remove(&pid);
                            }
                            for &(_, ref child) in children.values() {
                                child.signal(SIGTERM).ok();
                            }
                            break 'signal_loop;
                        }
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                _ => unreachable!(),
            }
        }
    }

    // Stopping loop
    if children.len() > 0 {
        let timeo = sup.kill_unresponsive_after;
        let mut deadline = SteadyTime::now() + Duration::seconds(timeo as i64);
        loop {
            match trap.wait(deadline) {
                Some(SIGINT) => {}
                Some(SIGTERM) => {
                    for &(_, ref child) in children.values() {
                        child.signal(SIGTERM).ok();
                    }
                }
                Some(SIGCHLD) => {
                    for (pid, _) in reap_zombies() {
                        children.remove(&pid);
                    }
                    try!(tty_guard.check().map_err(|e|
                        format!("Error handling tty: {}", e)));
                    if children.len() == 0 {
                        break;
                    }
                }
                Some(_) => unreachable!(),
                None => {
                    println!(
                        "---------- \
                        Processes {:?} are still alive. Killing ... \
                        -----------\n\
                        To prevent killing of the processes set \
                        ``kill-unresponsive-after`` to some larger value",
                        children.values().map(|&(name, _)| name)
                            .collect::<Vec<_>>());
                    for &(_, ref child) in children.values() {
                        child.signal(SIGKILL).ok();
                    }
                    // Basically this deadline should never happen
                    deadline = SteadyTime::now() + Duration::seconds(3600);
                }
            }
        }
    }

    if errcode == 0 {
        if let Some(ref epilog) = sup.epilog {
            print!("{}", epilog);
        }
    }

    Ok(errcode)
}
<|endoftext|>


<commit_before>//! Helper to run a callback in the libuv main thread.

use std::mem;
use std::os::raw::c_void;

use types::*;
use handle::{Handle, Managed};
use neon_runtime;
use neon_runtime::raw;
use std::sync::Arc;
use context::Context;

type EventContext<'a> = crate::context::TaskContext<'a>;

struct EventHandlerInner(*mut c_void);

unsafe impl Send for EventHandlerInner {}
unsafe impl Sync for EventHandlerInner {}

impl Drop for EventHandlerInner {
    fn drop(&mut self) {
        unsafe {
            neon_runtime::handler::delete(self.0);
        }
    }
}

#[derive(Clone)]
pub struct EventHandler(Arc<EventHandlerInner>);

impl EventHandler {
    #[cfg(feature = "napi-runtime")]
    pub fn new<'a, C: Context<'a>, T: Value>(cx: &C, this: Handle<T>, callback: Handle<JsFunction>) -> Self {
        unimplemented!()
    }

    #[cfg(feature = "legacy-runtime")]
    pub fn new<'a, C: Context<'a>, T: Value>(cx: &C, this: Handle<T>, callback: Handle<JsFunction>) -> Self {
        let cb = unsafe {
            neon_runtime::handler::new(cx.env().to_raw(), this.to_raw(), callback.to_raw())
        };
        EventHandler(Arc::new(EventHandlerInner(cb)))
    }

    pub fn schedule<T, F>(&self, arg_cb: F)
        where T: Value,
              F: for<'a> FnOnce(&mut EventContext<'a>) -> Vec<Handle<'a, T>>,
              F: Send + 'static {
        self.schedule_with(move |cx, this, callback| {
            let args = arg_cb(cx);
            let _result = callback.call(cx, this, args);
        })
    }

    pub fn schedule_with<F>(&self, arg_cb: F)
        where F: FnOnce(&mut EventContext, Handle<JsValue>, Handle<JsFunction>),
              F: Send + 'static {
        let callback = Box::into_raw(Box::new(arg_cb)) as *mut c_void;
        unsafe {
            neon_runtime::handler::schedule((*self.0).0, callback, handle_callback::<F>);
        }
    }
}

unsafe extern "C" fn handle_callback<F>(this: raw::Local, func: raw::Local, callback: *mut c_void)
    where F: FnOnce(&mut EventContext, Handle<JsValue>, Handle<JsFunction>), F: Send + 'static {
    EventContext::with(|mut cx: EventContext| {
        let this = JsValue::new_internal(this);
        let func: Handle<JsFunction> = Handle::new_internal(JsFunction::from_raw(func));
        let callback: Box<F> = Box::from_raw(mem::transmute(callback));
        callback(&mut cx, this, func);
    })
}
<commit_msg>Workaround for EventHandler drop race<commit_after>//! Helper to run a callback in the libuv main thread.

use std::mem;
use std::os::raw::c_void;

use types::*;
use handle::{Handle, Managed};
use neon_runtime;
use neon_runtime::raw;
use std::sync::Arc;
use context::Context;

type EventContext<'a> = crate::context::TaskContext<'a>;

struct EventHandlerInner(*mut c_void);

unsafe impl Send for EventHandlerInner {}
unsafe impl Sync for EventHandlerInner {}

impl Drop for EventHandlerInner {
    fn drop(&mut self) {
        unsafe {
            neon_runtime::handler::delete(self.0);
        }
    }
}

#[derive(Clone)]
pub struct EventHandler(Arc<EventHandlerInner>);

impl EventHandler {
    #[cfg(feature = "napi-runtime")]
    pub fn new<'a, C: Context<'a>, T: Value>(cx: &C, this: Handle<T>, callback: Handle<JsFunction>) -> Self {
        unimplemented!()
    }

    #[cfg(feature = "legacy-runtime")]
    pub fn new<'a, C: Context<'a>, T: Value>(cx: &C, this: Handle<T>, callback: Handle<JsFunction>) -> Self {
        let cb = unsafe {
            neon_runtime::handler::new(cx.env().to_raw(), this.to_raw(), callback.to_raw())
        };
        EventHandler(Arc::new(EventHandlerInner(cb)))
    }

    pub fn schedule<T, F>(&self, arg_cb: F)
        where T: Value,
              F: for<'a> FnOnce(&mut EventContext<'a>) -> Vec<Handle<'a, T>>,
              F: Send + 'static {
        self.schedule_with(move |cx, this, callback| {
            let args = arg_cb(cx);
            let _result = callback.call(cx, this, args);
        })
    }

    pub fn schedule_with<F>(&self, arg_cb: F)
        where F: FnOnce(&mut EventContext, Handle<JsValue>, Handle<JsFunction>),
              F: Send + 'static {
        // HACK: Work around for race condition in `close`. `EventHandler` cannot be
        // dropped until all callbacks have executed.
        // NOTE: This will still leak memory if the callback is never called
        let cloned_cb = self.clone();
        let cb = move |cx, this, cb| {
            arg_cb(cx, this, cb);
            std::mem::drop(cloned_cb);
        };

        let callback = Box::into_raw(Box::new(cb)) as *mut c_void;
        unsafe {
            neon_runtime::handler::schedule((*self.0).0, callback, handle_callback::<F>);
        }
    }
}

unsafe extern "C" fn handle_callback<F>(this: raw::Local, func: raw::Local, callback: *mut c_void)
    where F: FnOnce(&mut EventContext, Handle<JsValue>, Handle<JsFunction>), F: Send + 'static {
    EventContext::with(|mut cx: EventContext| {
        let this = JsValue::new_internal(this);
        let func: Handle<JsFunction> = Handle::new_internal(JsFunction::from_raw(func));
        let callback: Box<F> = Box::from_raw(mem::transmute(callback));
        callback(&mut cx, this, func);
    })
}
<|endoftext|>


<commit_before>//! Functions for computing gradients of image intensities.

use image::{
    GenericImage,
    ImageBuffer,
    Luma
};

use definitions::{
    VecBuffer
};

use filter::{
    filter3x3
};

/// Sobel filter for vertical edges.
static VERTICAL_SOBEL: [i32; 9] = [
    -1, 0, 1,
    -2, 0, 2,
    -1, 0, 1];

/// Sobel filter for horizontal edges.
static HORIZONTAL_SOBEL: [i32; 9] = [
    -1, -2, -1,
     0,  0,  0,
     1,  2,  1];

/// Convolves with the horizontal Sobel kernel to detect horizontal
/// edges in an image.
pub fn horizontal_sobel<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &HORIZONTAL_SOBEL)
}

/// Convolves with the vertical Sobel kernel to detect vertical
/// edges in an image.
pub fn vertical_sobel<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &VERTICAL_SOBEL)
}

/// Prewitt filter for vertical edges.
static VERTICAL_PREWITT: [i32; 9] = [
    -1, 0, 1,
    -1, 0, 1,
    -1, 0, 1];

/// Prewitt filter for horizontal edges.
static HORIZONTAL_PREWITT: [i32; 9] = [
    -1, -1, -1,
     0,  0,  0,
     1,  1,  1];

/// Convolves with the horizontal Prewitt kernel to detect horizontal
/// edges in an image.
pub fn horizontal_prewitt<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &HORIZONTAL_PREWITT)
}

/// Convolves with the vertical Prewitt kernel to detect vertical
/// edges in an image.
pub fn vertical_prewitt<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &VERTICAL_PREWITT)
}
/// Returns the magnitudes of gradients in an image.
// TODO: Returns directions as well as magnitudes.
// TODO: Support filtering with allocating a fresh image - filtering functions could
// TODO: take some kind of pixel-sink. This would allow us to compute gradient magnitudes
// TODO: and directions without allocating intermediates for vertical and horizontal gradients.
pub fn sobel_gradients<I>(image: &I) -> VecBuffer<Luma<u16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static {

    let horizontal = horizontal_sobel(image);
    let vertical = vertical_sobel(image);

    let (width, height) = image.dimensions();
    let mut out = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let h = horizontal.get_pixel(x, y)[0] as f32;
            let v = vertical.get_pixel(x, y)[0] as f32;
            let m = (h.powi(2) + v.powi(2)).sqrt() as u16;
            out.put_pixel(x, y, Luma([m]));
        }
    }

    out
}

#[cfg(test)]
mod test {

    use super::{
        horizontal_sobel,
        vertical_sobel,
        horizontal_prewitt,
        vertical_prewitt
    };

    use image::{
        GrayImage,
        ImageBuffer,
        Luma
    };

    #[test]
    fn test_gradients_constant_image() {
        let image = ImageBuffer::from_fn(5, 5, |_, _| Luma([15u8]));
        let expected = ImageBuffer::from_fn(5, 5, |_, _| Luma([0i16]));
        assert_pixels_eq!(horizontal_sobel(&image), expected);
        assert_pixels_eq!(vertical_sobel(&image), expected);
        assert_pixels_eq!(horizontal_prewitt(&image), expected);
        assert_pixels_eq!(vertical_prewitt(&image), expected);
    }

    #[test]
    fn test_vertical_sobel_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 2, 1,
            6, 5, 4,
            9, 8, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -4i16, -8i16, -4i16,
            -4i16, -8i16, -4i16,
            -4i16, -8i16, -4i16]).unwrap();

        let filtered = vertical_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_horizontal_sobel_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 6, 9,
            2, 5, 8,
            1, 4, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -4i16, -4i16, -4i16,
            -8i16, -8i16, -8i16,
            -4i16, -4i16, -4i16]).unwrap();

        let filtered = horizontal_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_vertical_prewitt_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 2, 1,
            6, 5, 4,
            9, 8, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -4i16, -8i16, -4i16,
            -4i16, -8i16, -4i16,
            -4i16, -8i16, -4i16]).unwrap();

        let filtered = vertical_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_horizontal_prewitt_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 6, 9,
            2, 5, 8,
            1, 4, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -4i16, -4i16, -4i16,
            -8i16, -8i16, -8i16,
            -4i16, -4i16, -4i16]).unwrap();

        let filtered = horizontal_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }
}
<commit_msg>Fix test<commit_after>//! Functions for computing gradients of image intensities.

use image::{
    GenericImage,
    ImageBuffer,
    Luma
};

use definitions::{
    VecBuffer
};

use filter::{
    filter3x3
};

/// Sobel filter for vertical edges.
static VERTICAL_SOBEL: [i32; 9] = [
    -1, 0, 1,
    -2, 0, 2,
    -1, 0, 1];

/// Sobel filter for horizontal edges.
static HORIZONTAL_SOBEL: [i32; 9] = [
    -1, -2, -1,
     0,  0,  0,
     1,  2,  1];

/// Convolves with the horizontal Sobel kernel to detect horizontal
/// edges in an image.
pub fn horizontal_sobel<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &HORIZONTAL_SOBEL)
}

/// Convolves with the vertical Sobel kernel to detect vertical
/// edges in an image.
pub fn vertical_sobel<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &VERTICAL_SOBEL)
}

/// Prewitt filter for vertical edges.
static VERTICAL_PREWITT: [i32; 9] = [
    -1, 0, 1,
    -1, 0, 1,
    -1, 0, 1];

/// Prewitt filter for horizontal edges.
static HORIZONTAL_PREWITT: [i32; 9] = [
    -1, -1, -1,
     0,  0,  0,
     1,  1,  1];

/// Convolves with the horizontal Prewitt kernel to detect horizontal
/// edges in an image.
pub fn horizontal_prewitt<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &HORIZONTAL_PREWITT)
}

/// Convolves with the vertical Prewitt kernel to detect vertical
/// edges in an image.
pub fn vertical_prewitt<I>(image: &I) -> VecBuffer<Luma<i16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static
{
    filter3x3(image, &VERTICAL_PREWITT)
}
/// Returns the magnitudes of gradients in an image.
// TODO: Returns directions as well as magnitudes.
// TODO: Support filtering with allocating a fresh image - filtering functions could
// TODO: take some kind of pixel-sink. This would allow us to compute gradient magnitudes
// TODO: and directions without allocating intermediates for vertical and horizontal gradients.
pub fn sobel_gradients<I>(image: &I) -> VecBuffer<Luma<u16>>
    where I: GenericImage<Pixel=Luma<u8>> + 'static {

    let horizontal = horizontal_sobel(image);
    let vertical = vertical_sobel(image);

    let (width, height) = image.dimensions();
    let mut out = ImageBuffer::new(width, height);

    for y in 0..height {
        for x in 0..width {
            let h = horizontal.get_pixel(x, y)[0] as f32;
            let v = vertical.get_pixel(x, y)[0] as f32;
            let m = (h.powi(2) + v.powi(2)).sqrt() as u16;
            out.put_pixel(x, y, Luma([m]));
        }
    }

    out
}

#[cfg(test)]
mod test {

    use super::{
        horizontal_sobel,
        vertical_sobel,
        horizontal_prewitt,
        vertical_prewitt
    };

    use image::{
        GrayImage,
        ImageBuffer,
        Luma
    };

    #[test]
    fn test_gradients_constant_image() {
        let image = ImageBuffer::from_fn(5, 5, |_, _| Luma([15u8]));
        let expected = ImageBuffer::from_fn(5, 5, |_, _| Luma([0i16]));
        assert_pixels_eq!(horizontal_sobel(&image), expected);
        assert_pixels_eq!(vertical_sobel(&image), expected);
        assert_pixels_eq!(horizontal_prewitt(&image), expected);
        assert_pixels_eq!(vertical_prewitt(&image), expected);
    }

    #[test]
    fn test_vertical_sobel_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 2, 1,
            6, 5, 4,
            9, 8, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -4i16, -8i16, -4i16,
            -4i16, -8i16, -4i16,
            -4i16, -8i16, -4i16]).unwrap();

        let filtered = vertical_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_horizontal_sobel_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 6, 9,
            2, 5, 8,
            1, 4, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -4i16, -4i16, -4i16,
            -8i16, -8i16, -8i16,
            -4i16, -4i16, -4i16]).unwrap();

        let filtered = horizontal_sobel(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_vertical_prewitt_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 2, 1,
            6, 5, 4,
            9, 8, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -3i16, -6i16, -3i16,
            -3i16, -6i16, -3i16,
            -3i16, -6i16, -3i16]).unwrap();

        let filtered = vertical_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }

    #[test]
    fn test_horizontal_prewitt_gradient_image() {
        let image: GrayImage = ImageBuffer::from_raw(3, 3, vec![
            3, 6, 9,
            2, 5, 8,
            1, 4, 7]).unwrap();

        let expected = ImageBuffer::from_raw(3, 3, vec![
            -3i16, -3i16, -3i16,
            -6i16, -6i16, -6i16,
            -3i16, -3i16, -3i16]).unwrap();

        let filtered = horizontal_prewitt(&image);
        assert_pixels_eq!(filtered, expected);
    }
}
<|endoftext|>


<commit_before>// Std
use std::{collections::HashMap, ffi::OsString, mem, ops::Deref};

// Internal
use crate::{
    build::{App, Arg, ArgSettings},
    parse::{ArgMatches, MatchedArg, SubCommand, ValueType},
    util::Id,
};

// Third party
use indexmap::map::Entry;

#[derive(Debug, Default)]
pub(crate) struct ArgMatcher(pub(crate) ArgMatches);

impl ArgMatcher {
    pub(crate) fn new(_app: &App) -> Self {
        ArgMatcher(ArgMatches {
            #[cfg(debug_assertions)]
            valid_args: {
                let args = _app.args.args().map(|a| a.id.clone());
                let groups = _app.groups.iter().map(|g| g.id.clone());
                args.chain(groups).collect()
            },
            #[cfg(debug_assertions)]
            valid_subcommands: _app.subcommands.iter().map(|sc| sc.id.clone()).collect(),
            // HACK: Allow an external subcommand's ArgMatches be a stand-in for any ArgMatches
            // since users can't detect it and avoid the asserts.
            //
            // See clap-rs/clap#3263
            #[cfg(debug_assertions)]
            disable_asserts: _app.is_set(crate::AppSettings::AllowExternalSubcommands),
            ..Default::default()
        })
    }

    pub(crate) fn into_inner(self) -> ArgMatches {
        self.0
    }

    pub(crate) fn propagate_globals(&mut self, global_arg_vec: &[Id]) {
        debug!(
            "ArgMatcher::get_global_values: global_arg_vec={:?}",
            global_arg_vec
        );
        let mut vals_map = HashMap::new();
        self.fill_in_global_values(global_arg_vec, &mut vals_map);
    }

    fn fill_in_global_values(
        &mut self,
        global_arg_vec: &[Id],
        vals_map: &mut HashMap<Id, MatchedArg>,
    ) {
        for global_arg in global_arg_vec {
            if let Some(ma) = self.get(global_arg) {
                // We have to check if the parent's global arg wasn't used but still exists
                // such as from a default value.
                //
                // For example, `myprog subcommand --global-arg=value` where --global-arg defines
                // a default value of `other` myprog would have an existing MatchedArg for
                // --global-arg where the value is `other`, however the occurs will be 0.
                let to_update = if let Some(parent_ma) = vals_map.get(global_arg) {
                    if parent_ma.occurs > 0 && ma.occurs == 0 {
                        parent_ma.clone()
                    } else {
                        ma.clone()
                    }
                } else {
                    ma.clone()
                };
                vals_map.insert(global_arg.clone(), to_update);
            }
        }
        if let Some(ref mut sc) = self.0.subcommand {
            let mut am = ArgMatcher(mem::take(&mut sc.matches));
            am.fill_in_global_values(global_arg_vec, vals_map);
            mem::swap(&mut am.0, &mut sc.matches);
        }

        for (name, matched_arg) in vals_map.iter_mut() {
            self.0.args.insert(name.clone(), matched_arg.clone());
        }
    }

    pub(crate) fn get_mut(&mut self, arg: &Id) -> Option<&mut MatchedArg> {
        self.0.args.get_mut(arg)
    }

    pub(crate) fn get(&self, arg: &Id) -> Option<&MatchedArg> {
        self.0.args.get(arg)
    }

    pub(crate) fn remove(&mut self, arg: &Id) {
        self.0.args.swap_remove(arg);
    }

    pub(crate) fn contains(&self, arg: &Id) -> bool {
        self.0.args.contains_key(arg)
    }

    pub(crate) fn contains_explicit(&self, arg: &Id) -> bool {
        self.0
            .args
            .get(arg)
            .map_or(false, |a| a.ty != ValueType::DefaultValue)
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.0.args.is_empty()
    }

    pub(crate) fn arg_names(&self) -> indexmap::map::Keys<Id, MatchedArg> {
        self.0.args.keys()
    }

    pub(crate) fn entry(&mut self, arg: &Id) -> indexmap::map::Entry<Id, MatchedArg> {
        self.0.args.entry(arg.clone())
    }

    pub(crate) fn subcommand(&mut self, sc: SubCommand) {
        self.0.subcommand = Some(Box::new(sc));
    }

    pub(crate) fn subcommand_name(&self) -> Option<&str> {
        self.0.subcommand_name()
    }

    pub(crate) fn iter(&self) -> indexmap::map::Iter<Id, MatchedArg> {
        self.0.args.iter()
    }

    pub(crate) fn inc_occurrence_of_arg(&mut self, arg: &Arg) {
        let id = &arg.id;
        debug!("ArgMatcher::inc_occurrence_of_arg: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new());
        ma.update_ty(ValueType::CommandLine);
        ma.set_ignore_case(arg.is_set(ArgSettings::IgnoreCase));
        ma.invalid_utf8_allowed(arg.is_set(ArgSettings::AllowInvalidUtf8));
        ma.occurs += 1;
    }

    pub(crate) fn inc_occurrence_of_group(&mut self, id: &Id) {
        debug!("ArgMatcher::inc_occurrence_of_group: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new());
        ma.update_ty(ValueType::CommandLine);
        ma.occurs += 1;
    }

    pub(crate) fn add_val_to(&mut self, arg: &Id, val: OsString, ty: ValueType, append: bool) {
        if append {
            self.append_val_to(arg, val, ty);
        } else {
            self.push_val_to(arg, val, ty);
        }
    }

    fn push_val_to(&mut self, arg: &Id, val: OsString, ty: ValueType) {
        // We will manually inc occurrences later(for flexibility under
        // specific circumstances, like only add one occurrence for flag
        // when we met: `--flag=one,two`).
        let ma = self.entry(arg).or_default();
        ma.update_ty(ty);
        ma.push_val(val);
    }

    fn append_val_to(&mut self, arg: &Id, val: OsString, ty: ValueType) {
        let ma = self.entry(arg).or_default();
        ma.update_ty(ty);
        ma.append_val(val);
    }

    pub(crate) fn new_val_group(&mut self, arg: &Id) {
        let ma = self.entry(arg).or_default();
        ma.new_val_group();
    }

    pub(crate) fn add_index_to(&mut self, arg: &Id, idx: usize, ty: ValueType) {
        let ma = self.entry(arg).or_default();
        ma.update_ty(ty);
        ma.push_index(idx);
    }

    pub(crate) fn has_val_groups(&mut self, arg: &Id) -> bool {
        match self.entry(arg) {
            Entry::Occupied(e) => e.get().has_val_groups(),
            Entry::Vacant(_) => false,
        }
    }

    pub(crate) fn needs_more_vals(&self, o: &Arg) -> bool {
        debug!("ArgMatcher::needs_more_vals: o={}", o.name);
        if let Some(ma) = self.get(&o.id) {
            let current_num = ma.num_vals();
            if let Some(num) = o.num_vals {
                debug!("ArgMatcher::needs_more_vals: num_vals...{}", num);
                return if o.is_set(ArgSettings::MultipleOccurrences) {
                    (current_num % num) != 0
                } else {
                    num != current_num
                };
            } else if let Some(num) = o.max_vals {
                debug!("ArgMatcher::needs_more_vals: max_vals...{}", num);
                return current_num < num;
            } else if o.min_vals.is_some() {
                debug!("ArgMatcher::needs_more_vals: min_vals...true");
                return true;
            }
            return o.is_set(ArgSettings::MultipleValues);
        }
        true
    }
}

impl Deref for ArgMatcher {
    type Target = ArgMatches;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
<commit_msg>perf: Merge .clone() calls<commit_after>// Std
use std::{collections::HashMap, ffi::OsString, mem, ops::Deref};

// Internal
use crate::{
    build::{App, Arg, ArgSettings},
    parse::{ArgMatches, MatchedArg, SubCommand, ValueType},
    util::Id,
};

// Third party
use indexmap::map::Entry;

#[derive(Debug, Default)]
pub(crate) struct ArgMatcher(pub(crate) ArgMatches);

impl ArgMatcher {
    pub(crate) fn new(_app: &App) -> Self {
        ArgMatcher(ArgMatches {
            #[cfg(debug_assertions)]
            valid_args: {
                let args = _app.args.args().map(|a| a.id.clone());
                let groups = _app.groups.iter().map(|g| g.id.clone());
                args.chain(groups).collect()
            },
            #[cfg(debug_assertions)]
            valid_subcommands: _app.subcommands.iter().map(|sc| sc.id.clone()).collect(),
            // HACK: Allow an external subcommand's ArgMatches be a stand-in for any ArgMatches
            // since users can't detect it and avoid the asserts.
            //
            // See clap-rs/clap#3263
            #[cfg(debug_assertions)]
            disable_asserts: _app.is_set(crate::AppSettings::AllowExternalSubcommands),
            ..Default::default()
        })
    }

    pub(crate) fn into_inner(self) -> ArgMatches {
        self.0
    }

    pub(crate) fn propagate_globals(&mut self, global_arg_vec: &[Id]) {
        debug!(
            "ArgMatcher::get_global_values: global_arg_vec={:?}",
            global_arg_vec
        );
        let mut vals_map = HashMap::new();
        self.fill_in_global_values(global_arg_vec, &mut vals_map);
    }

    fn fill_in_global_values(
        &mut self,
        global_arg_vec: &[Id],
        vals_map: &mut HashMap<Id, MatchedArg>,
    ) {
        for global_arg in global_arg_vec {
            if let Some(ma) = self.get(global_arg) {
                // We have to check if the parent's global arg wasn't used but still exists
                // such as from a default value.
                //
                // For example, `myprog subcommand --global-arg=value` where --global-arg defines
                // a default value of `other` myprog would have an existing MatchedArg for
                // --global-arg where the value is `other`, however the occurs will be 0.
                let to_update = if let Some(parent_ma) = vals_map.get(global_arg) {
                    if parent_ma.occurs > 0 && ma.occurs == 0 {
                        parent_ma
                    } else {
                        ma
                    }
                } else {
                    ma
                }
                .clone();
                vals_map.insert(global_arg.clone(), to_update);
            }
        }
        if let Some(ref mut sc) = self.0.subcommand {
            let mut am = ArgMatcher(mem::take(&mut sc.matches));
            am.fill_in_global_values(global_arg_vec, vals_map);
            mem::swap(&mut am.0, &mut sc.matches);
        }

        for (name, matched_arg) in vals_map.iter_mut() {
            self.0.args.insert(name.clone(), matched_arg.clone());
        }
    }

    pub(crate) fn get_mut(&mut self, arg: &Id) -> Option<&mut MatchedArg> {
        self.0.args.get_mut(arg)
    }

    pub(crate) fn get(&self, arg: &Id) -> Option<&MatchedArg> {
        self.0.args.get(arg)
    }

    pub(crate) fn remove(&mut self, arg: &Id) {
        self.0.args.swap_remove(arg);
    }

    pub(crate) fn contains(&self, arg: &Id) -> bool {
        self.0.args.contains_key(arg)
    }

    pub(crate) fn contains_explicit(&self, arg: &Id) -> bool {
        self.0
            .args
            .get(arg)
            .map_or(false, |a| a.ty != ValueType::DefaultValue)
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.0.args.is_empty()
    }

    pub(crate) fn arg_names(&self) -> indexmap::map::Keys<Id, MatchedArg> {
        self.0.args.keys()
    }

    pub(crate) fn entry(&mut self, arg: &Id) -> indexmap::map::Entry<Id, MatchedArg> {
        self.0.args.entry(arg.clone())
    }

    pub(crate) fn subcommand(&mut self, sc: SubCommand) {
        self.0.subcommand = Some(Box::new(sc));
    }

    pub(crate) fn subcommand_name(&self) -> Option<&str> {
        self.0.subcommand_name()
    }

    pub(crate) fn iter(&self) -> indexmap::map::Iter<Id, MatchedArg> {
        self.0.args.iter()
    }

    pub(crate) fn inc_occurrence_of_arg(&mut self, arg: &Arg) {
        let id = &arg.id;
        debug!("ArgMatcher::inc_occurrence_of_arg: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new());
        ma.update_ty(ValueType::CommandLine);
        ma.set_ignore_case(arg.is_set(ArgSettings::IgnoreCase));
        ma.invalid_utf8_allowed(arg.is_set(ArgSettings::AllowInvalidUtf8));
        ma.occurs += 1;
    }

    pub(crate) fn inc_occurrence_of_group(&mut self, id: &Id) {
        debug!("ArgMatcher::inc_occurrence_of_group: id={:?}", id);
        let ma = self.entry(id).or_insert(MatchedArg::new());
        ma.update_ty(ValueType::CommandLine);
        ma.occurs += 1;
    }

    pub(crate) fn add_val_to(&mut self, arg: &Id, val: OsString, ty: ValueType, append: bool) {
        if append {
            self.append_val_to(arg, val, ty);
        } else {
            self.push_val_to(arg, val, ty);
        }
    }

    fn push_val_to(&mut self, arg: &Id, val: OsString, ty: ValueType) {
        // We will manually inc occurrences later(for flexibility under
        // specific circumstances, like only add one occurrence for flag
        // when we met: `--flag=one,two`).
        let ma = self.entry(arg).or_default();
        ma.update_ty(ty);
        ma.push_val(val);
    }

    fn append_val_to(&mut self, arg: &Id, val: OsString, ty: ValueType) {
        let ma = self.entry(arg).or_default();
        ma.update_ty(ty);
        ma.append_val(val);
    }

    pub(crate) fn new_val_group(&mut self, arg: &Id) {
        let ma = self.entry(arg).or_default();
        ma.new_val_group();
    }

    pub(crate) fn add_index_to(&mut self, arg: &Id, idx: usize, ty: ValueType) {
        let ma = self.entry(arg).or_default();
        ma.update_ty(ty);
        ma.push_index(idx);
    }

    pub(crate) fn has_val_groups(&mut self, arg: &Id) -> bool {
        match self.entry(arg) {
            Entry::Occupied(e) => e.get().has_val_groups(),
            Entry::Vacant(_) => false,
        }
    }

    pub(crate) fn needs_more_vals(&self, o: &Arg) -> bool {
        debug!("ArgMatcher::needs_more_vals: o={}", o.name);
        if let Some(ma) = self.get(&o.id) {
            let current_num = ma.num_vals();
            if let Some(num) = o.num_vals {
                debug!("ArgMatcher::needs_more_vals: num_vals...{}", num);
                return if o.is_set(ArgSettings::MultipleOccurrences) {
                    (current_num % num) != 0
                } else {
                    num != current_num
                };
            } else if let Some(num) = o.max_vals {
                debug!("ArgMatcher::needs_more_vals: max_vals...{}", num);
                return current_num < num;
            } else if o.min_vals.is_some() {
                debug!("ArgMatcher::needs_more_vals: min_vals...true");
                return true;
            }
            return o.is_set(ArgSettings::MultipleValues);
        }
        true
    }
}

impl Deref for ArgMatcher {
    type Target = ArgMatches;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}
<|endoftext|>