Coverage for harbor_cli/commands/cli/find.py: 98%

59 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-09 12:09 +0100

1from __future__ import annotations 

2 

3from enum import Enum 

4from typing import Callable 

5from typing import List 

6from typing import Optional 

7 

8import typer 

9from fuzzywuzzy import fuzz 

10 

11from ...app import app 

12from ...logs import logger 

13from ...models import CommandSummary 

14from ...output.render import render_result # avoid circular import # TODO: fix 

15from ...utils.commands import get_app_commands 

16 

17 

18class MatchStrategy(Enum): 

19 """Match strategies.""" 

20 

21 RATIO = "ratio" 

22 PARTIAL_RATIO = "partial-ratio" 

23 TOKEN_SORT_RATIO = "token-sort-ratio" 

24 TOKEN_SET_RATIO = "token-set-ratio" 

25 

26 

27def get_match_func(strategy: MatchStrategy) -> Callable[[str, str], int]: 

28 """Get the match function for a strategy.""" 

29 if strategy == MatchStrategy.RATIO: 

30 return fuzz.ratio 

31 elif strategy == MatchStrategy.PARTIAL_RATIO: 

32 return fuzz.partial_ratio 

33 elif strategy == MatchStrategy.TOKEN_SORT_RATIO: 

34 return fuzz.token_sort_ratio 

35 elif strategy == MatchStrategy.TOKEN_SET_RATIO: 

36 return fuzz.token_set_ratio 

37 else: 

38 raise ValueError(f"Unknown match strategy: {strategy}") # pragma: no cover 

39 

40 

41def match_commands( 

42 commands: list[CommandSummary], 

43 query: str, 

44 min_score: int, 

45 names: bool, 

46 descriptions: bool, 

47 strategy: MatchStrategy, 

48) -> list[CommandSummary]: 

49 """Use fuzzy matching to find commands that match a query.""" 

50 match_func = get_match_func(strategy) 

51 

52 matches = [] 

53 for command in commands: 

54 score = 0 

55 if names: 

56 name_score = match_func(query, command.name) 

57 if name_score > score: 

58 score = name_score 

59 if descriptions: 

60 desc_score = match_func(query, command.help) 

61 if desc_score > score: 

62 score = desc_score 

63 command.score = score 

64 if command.score >= min_score: 

65 matches.append(command) 

66 return sorted(matches, key=lambda x: x.score, reverse=True) 

67 

68 

69@app.command(no_args_is_help=True, name="find") 

70def find_command( 

71 ctx: typer.Context, 

72 query: List[str] = typer.Argument( 

73 ..., 

74 help="The search query.", 

75 ), 

76 limit: Optional[int] = typer.Option( 

77 None, 

78 help="Maximum number of results to show.", 

79 min=1, 

80 ), 

81 min_score: int = typer.Option( 

82 75, 

83 "--min-score", 

84 help="Minimum match ratio to show. Lower = more results.", 

85 min=0, 

86 max=100, 

87 ), 

88 names: bool = typer.Option( 

89 True, 

90 "--names/--no-names", 

91 help="Search in command names.", 

92 ), 

93 descriptions: bool = typer.Option( 

94 True, 

95 "--descriptions/--no-descriptions", 

96 help="Search in command descriptions.", 

97 ), 

98 strategy: MatchStrategy = typer.Option( 

99 MatchStrategy.PARTIAL_RATIO.value, 

100 "--strategy", 

101 help=( 

102 "The matching strategy to use. " 

103 "Strategies require different scoring thresholds. " 

104 "The default threshold is optimized for partial ratio." 

105 ), 

106 ), 

107) -> None: 

108 """Search commands based on names and descriptions.""" 

109 matches = _do_find( 

110 ctx=ctx, 

111 query=query, 

112 limit=limit, 

113 min_score=min_score, 

114 names=names, 

115 descriptions=descriptions, 

116 strategy=strategy, 

117 ) 

118 render_result(matches, ctx) 

119 

120 

121def _do_find( 

122 ctx: typer.Context, 

123 query: List[str], 

124 limit: Optional[int], 

125 min_score: int, 

126 names: bool, 

127 descriptions: bool, 

128 strategy: MatchStrategy, 

129) -> list[CommandSummary]: 

130 # Join arguments to a single string 

131 q = " ".join(query) 

132 

133 commands = get_app_commands(app) 

134 matches = match_commands( 

135 commands, 

136 q, 

137 min_score=min_score, 

138 names=names, 

139 descriptions=descriptions, 

140 strategy=strategy, 

141 ) 

142 if limit is not None and len(matches) > limit: 

143 logger.warning(f"{len(matches) - limit} results omitted.") 

144 matches = matches[:limit] 

145 return matches 

146 

147 

148@app.command(name="commands") 

149def list_commands(ctx: typer.Context) -> None: 

150 """List all commands.""" 

151 commands = get_app_commands(app) 

152 render_result(commands, ctx)