ohmyzsh/plugins/scd/scd

#!/bin/zsh -f

emulate -L zsh
if [[ $(whence -w $0) == *:' 'command ]]; then
    emulate -R zsh
    alias return=exit
    local RUNNING_AS_COMMAND=1
fi

local DOC='scd -- smart change to a recently used directory
usage: scd [options] [pattern1 pattern2 ...]
Go to a directory path that contains all fixed string patterns.  Prefer
recently visited directories and directories with patterns in their tail
component.  Display a selection menu in case of multiple matches.

Options:
  -a, --add         add specified directories to the directory index
  --unindex         remove specified directories from the index
  -r, --recursive   apply options --add or --unindex recursively
  --alias=ALIAS     create alias for the current or specified directory and
                    store it in ~/.scdalias.zsh
  --unalias         remove ALIAS definition for the current or specified
                    directory from ~/.scdalias.zsh
  --list            show matching directories and exit
  -v, --verbose     display directory rank in the selection menu
  -h, --help        display this message and exit
'

local SCD_HISTFILE=${SCD_HISTFILE:-${HOME}/.scdhistory}
local SCD_HISTSIZE=${SCD_HISTSIZE:-5000}
local SCD_MENUSIZE=${SCD_MENUSIZE:-20}
local SCD_MEANLIFE=${SCD_MEANLIFE:-86400}
local SCD_THRESHOLD=${SCD_THRESHOLD:-0.005}
local SCD_SCRIPT=${RUNNING_AS_COMMAND:+$SCD_SCRIPT}
local SCD_ALIAS=~/.scdalias.zsh

local ICASE a d m p i tdir maxrank threshold
local opt_help opt_add opt_unindex opt_recursive opt_verbose
local opt_alias opt_unalias opt_list
local -A drank dalias dkey
local dmatching

setopt extendedhistory extendedglob noautonamedirs brace_ccl

# If SCD_SCRIPT is defined make sure the file exists and is empty.
# This removes any previous old commands.
[[ -n "$SCD_SCRIPT" ]] && [[ -s $SCD_SCRIPT || ! -f $SCD_SCRIPT ]] && (
    umask 077
    : >| $SCD_SCRIPT
)

# process command line options
zmodload -i zsh/zutil
zmodload -i zsh/datetime
zparseopts -D -- a=opt_add -add=opt_add -unindex=opt_unindex \
    r=opt_recursive -recursive=opt_recursive \
    -alias:=opt_alias -unalias=opt_unalias -list=opt_list \
    v=opt_verbose -verbose=opt_verbose h=opt_help -help=opt_help \
    || return $?

if [[ -n $opt_help ]]; then
    print $DOC
    return
fi

# load directory aliases if they exist
[[ -r $SCD_ALIAS ]] && source $SCD_ALIAS

# works faster than the (:a) modifier and is compatible with zsh 4.2.6
_scd_Y19oug_abspath() {
    set -A $1 ${(ps:\0:)"$(
        unfunction -m "*"; shift
        for d; do
            cd $d && print -Nr -- $PWD && cd $OLDPWD
        done
        )"}
}

# define directory alias
if [[ -n $opt_alias ]]; then
    if [[ -n $1 && ! -d $1 ]]; then
        print -u2 "'$1' is not a directory"
        return 1
    fi
    a=${opt_alias[-1]#=}
    _scd_Y19oug_abspath d ${1:-$PWD}
    # alias in the current shell, update alias file if successful
    hash -d -- $a=$d &&
    (
        umask 077
        hash -dr
        [[ -r $SCD_ALIAS ]] && source $SCD_ALIAS
        hash -d -- $a=$d
        hash -dL >| $SCD_ALIAS
    )
    return $?
fi

# undefine directory alias
if [[ -n $opt_unalias ]]; then
    if [[ -n $1 && ! -d $1 ]]; then
        print -u2 "'$1' is not a directory"
        return 1
    fi
    _scd_Y19oug_abspath a ${1:-$PWD}
    a=$(print -rD ${a})
    if [[ $a != [~][^/]## ]]; then
        return
    fi
    a=${a#[~]}
    # unalias in the current shell, update alias file if successful
    if unhash -d -- $a 2>/dev/null && [[ -r $SCD_ALIAS ]]; then
        (
            umask 077
            hash -dr
            source $SCD_ALIAS
            unhash -d -- $a 2>/dev/null &&
            hash -dL >| $SCD_ALIAS
        )
    fi
    return $?
fi

# Rewrite the history file if it is at least 20% oversized
if [[ -s $SCD_HISTFILE ]] && \
(( $(wc -l <$SCD_HISTFILE) > 1.2 * $SCD_HISTSIZE )); then
    m=( ${(f)"$(<$SCD_HISTFILE)"} )
    print -lr -- ${m[-$SCD_HISTSIZE,-1]} >| ${SCD_HISTFILE}
fi

# Internal functions are prefixed with "_scd_Y19oug_".
# The "record" function adds a non-repeating directory to the history
# and turns on history writing.
_scd_Y19oug_record() {
    while [[ -n $1 && $1 == ${history[$HISTCMD]} ]]; do
        shift
    done
    if [[ $# != 0 ]]; then
        ( umask 077; : >>| $SCD_HISTFILE )
        p=": ${EPOCHSECONDS}:0;"
        print -lr -- ${p}${^*} >> $SCD_HISTFILE
    fi
}

if [[ -n $opt_add ]]; then
    for a; do
        if [[ ! -d $a ]]; then
            print -u 2 "Directory $a does not exist"
            return 2
        fi
    done
    _scd_Y19oug_abspath m ${*:-$PWD}
    _scd_Y19oug_record $m
    if [[ -n $opt_recursive ]]; then
        for d in $m; do
            print -n "scanning ${d} ... "
            _scd_Y19oug_record ${d}/**/*(-/N)
            print "[done]"
        done
    fi
    return
fi

# take care of removing entries from the directory index
if [[ -n $opt_unindex ]]; then
    if [[ ! -s $SCD_HISTFILE ]]; then
        return
    fi
    # expand existing directories in the argument list
    for i in {1..$#}; do
        if [[ -d ${argv[i]} ]]; then
            _scd_Y19oug_abspath d ${argv[i]}
            argv[i]=${d}
        fi
    done
    m="$(awk -v recursive=${opt_recursive} '
        BEGIN {
            for (i = 2; i < ARGC; ++i) {
                argset[ARGV[i]] = 1;
                delete ARGV[i];
            }
        }
        1 {
            d = $0;  sub(/^[^;]*;/, "", d);
            if (d in argset)  next;
        }
        recursive {
            for (a in argset) {
                if (substr(d, 1, length(a) + 1) == a"/")  next;
            }
        }
        { print $0 }
        ' $SCD_HISTFILE ${*:-$PWD} )" || return $?
    : >| ${SCD_HISTFILE}
    [[ ${#m} == 0 ]] || print -r -- $m >> ${SCD_HISTFILE}
    return
fi

# The "action" function is called when there is just one target directory.
_scd_Y19oug_action() {
    if [[ -n $opt_list ]]; then
        for d; do
            a=${(k)dalias[(r)${d}]}
            print -r -- "# $a"
            print -r -- $d
        done
    elif [[ $# == 1 ]]; then
        if [[ -z $SCD_SCRIPT && -n $RUNNING_AS_COMMAND ]]; then
            print -u2 "Warning: running as command with SCD_SCRIPT undefined."
        fi
        [[ -n $SCD_SCRIPT ]] && (umask 077;
            print -r "cd ${(q)1}" >| $SCD_SCRIPT)
        [[ -N $SCD_HISTFILE ]] && touch -a $SCD_HISTFILE
        cd $1
        # record the new directory unless already done in some chpwd hook
        [[ -N $SCD_HISTFILE ]] || _scd_Y19oug_record $PWD
    fi
}

# handle different argument scenarios ----------------------------------------

## single argument that is an existing directory
if [[ $# == 1 && -d $1 && -x $1 ]]; then
    _scd_Y19oug_action $1
    return $?
## single argument that is an alias
elif [[ $# == 1 && -d ${d::=${nameddirs[$1]}} ]]; then
    _scd_Y19oug_action $d
    return $?
fi

# ignore case unless there is an argument with an uppercase letter
[[ "$*" == *[[:upper:]]* ]] || ICASE='(#i)'

# calculate rank of all directories in the SCD_HISTFILE and keep it as drank
# include a dummy entry for splitting of an empty string is buggy
[[ -s $SCD_HISTFILE ]] && drank=( ${(f)"$(
    print -l /dev/null -10
    <$SCD_HISTFILE \
    awk -v epochseconds=$EPOCHSECONDS -v meanlife=$SCD_MEANLIFE '
        BEGIN { FS = "[:;]"; }
        length($0) < 4096 && $2 > 0 {
            tau = 1.0 * ($2 - epochseconds) / meanlife;
            if (tau < -4.61)  tau = -4.61;
            prec = exp(tau);
            sub(/^[^;]*;/, "");
            if (NF)  ptot[$0] += prec;
        }
        END { for (di in ptot)  { print di; print ptot[di]; } }'
    )"}
)
unset "drank[/dev/null]"

# filter drank to the entries that match all arguments
for a; do
    p=${ICASE}"*${a}*"
    drank=( ${(kv)drank[(I)${~p}]} )
done

# build a list of matching directories reverse-sorted by their probabilities
dmatching=( ${(f)"$(
    for d p in ${(kv)drank}; do
        print -r -- "$p $d";
    done | sort -grk1 | cut -d ' ' -f 2-
    )"}
)

# if some directory paths match all patterns in order, discard all others
p=${ICASE}"*${(j:*:)argv}*"
m=( ${(M)dmatching:#${~p}} )
[[ -d ${m[1]} ]] && dmatching=( $m )
# if some directory names match last pattern, discard all others
p=${ICASE}"*${(j:*:)argv}[^/]#"
m=( ${(M)dmatching:#${~p}} )
[[ -d ${m[1]} ]] && dmatching=( $m )
# if some directory names match all patterns, discard all others
m=( $dmatching )
for a; do
    p=${ICASE}"*/[^/]#${a}[^/]#"
    m=( ${(M)m:#${~p}} )
done
[[ -d ${m[1]} ]] && dmatching=( $m )
# if some directory names match all patterns in order, discard all others
p=${ICASE}"/*${(j:[^/]#:)argv}[^/]#"
m=( ${(M)dmatching:#${~p}} )
[[ -d ${m[1]} ]] && dmatching=( $m )

# do not match $HOME or $PWD when run without arguments
if [[ $# == 0 ]]; then
    dmatching=( ${dmatching:#(${HOME}|${PWD})} )
fi

# keep at most SCD_MENUSIZE of matching and valid directories
m=( )
for d in $dmatching; do
    [[ ${#m} == $SCD_MENUSIZE ]] && break
    [[ -d $d && -x $d ]] && m+=$d
done
dmatching=( $m )

# find the maximum rank
maxrank=0.0
for d in $dmatching; do
    [[ ${drank[$d]} -lt maxrank ]] || maxrank=${drank[$d]}
done

# discard all directories below the rank threshold
threshold=$(( maxrank * SCD_THRESHOLD ))
dmatching=( ${^dmatching}(Ne:'(( ${drank[$REPLY]} >= threshold ))':) )

## process whatever directories that remained
case ${#dmatching} in
(0)
    print -u2 "no matching directory"
    return 1
    ;;
(1)
    _scd_Y19oug_action $dmatching
    return $?
    ;;
(*)
    # build a list of strings to be displayed in the selection menu
    m=( ${(f)"$(print -lD ${dmatching})"} )
    if [[ -n $opt_verbose ]]; then
        for i in {1..${#dmatching}}; do
            d=${dmatching[i]}
            m[i]=$(printf "%.3g %s" ${drank[$d]} $d)
        done
    fi
    # build a map of string names to actual directory paths
    for i in {1..${#m}}; dalias[${m[i]}]=${dmatching[i]}
    # opt_list - output matching directories and exit
    if [[ -n $opt_list ]]; then
        _scd_Y19oug_action ${dmatching}
        return
    fi
    # finally use the selection menu to get the answer
    a=( {a-z} {A-Z} )
    p=( )
    for i in {1..${#m}}; do
        [[ -n ${a[i]} ]] || break
        dkey[${a[i]}]=${dalias[$m[i]]}
        p+="${a[i]}) ${m[i]}"
    done
    print -c -r -- $p
    if read -s -k 1 d && [[ -n ${dkey[$d]} ]]; then
        _scd_Y19oug_action ${dkey[$d]}
    fi
    return $?
esac