Skip to content

Commit ef2435c

Browse files
kceiwkceiw
committed
Improve getting command ast and parse parameters.
- We use the RelatedAsts from PSReadLine to get the command AST as the user input. This is what our prediction will try to match. - Now we can parse the parameter in the format of "-Name:Value". - We don't parse positioinal parameters yet.
1 parent 1f77b70 commit ef2435c

File tree

9 files changed

+167
-66
lines changed

9 files changed

+167
-66
lines changed

‎tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorServiceTests.cs

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -89,10 +89,10 @@ public void VerifyParameterValues()
8989
Action actual = () => this._service.GetSuggestion(null, 1, 1, CancellationToken.None);
9090
Assert.Throws<ArgumentNullException>(actual);
9191

92-
actual = () => this._service.GetSuggestion(predictionContext.InputAst, 0, 1, CancellationToken.None);
92+
actual = () => this._service.GetSuggestion(predictionContext, 0, 1, CancellationToken.None);
9393
Assert.Throws<ArgumentOutOfRangeException>(actual);
9494

95-
actual = () => this._service.GetSuggestion(predictionContext.InputAst, 1, 0, CancellationToken.None);
95+
actual = () => this._service.GetSuggestion(predictionContext, 1, 0, CancellationToken.None);
9696
Assert.Throws<ArgumentOutOfRangeException>(actual);
9797
}
9898

@@ -110,8 +110,8 @@ public void VerifyParameterValues()
110110
public void VerifyUsingCommandBasedPredictor(string userInput)
111111
{
112112
var predictionContext = PredictionContext.Create(userInput);
113-
var commandAst = predictionContext.InputAst.FindAll(p => p is CommandAst, true).LastOrDefault() as CommandAst;
114-
var commandName = (commandAst?.CommandElements?.FirstOrDefault() as StringConstantExpressionAst)?.Value;
113+
var commandAst = predictionContext.RelatedAsts.OfType<CommandAst>().LastOrDefault();
114+
var commandName = commandAst?.GetCommandName();
115115
var inputParameterSet = new ParameterSet(commandAst);
116116
var rawUserInput = predictionContext.InputAst.Extent.Text;
117117
var presentCommands = new Dictionary<string, int>();
@@ -123,7 +123,7 @@ public void VerifyUsingCommandBasedPredictor(string userInput)
123123
1,
124124
CancellationToken.None);
125125

126-
var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
126+
var actual = this._service.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
127127
Assert.NotNull(actual);
128128
Assert.True(actual.Count > 0);
129129
Assert.NotNull(actual.PredictiveSuggestions.First());
@@ -133,7 +133,7 @@ public void VerifyUsingCommandBasedPredictor(string userInput)
133133
Assert.Equal<string>(expected.SourceTexts, actual.SourceTexts);
134134
Assert.All<SuggestionSource>(actual.SuggestionSources, (source) => Assert.Equal(SuggestionSource.CurrentCommand, source));
135135

136-
actual = this._noFallbackPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
136+
actual = this._noFallbackPredictorService.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
137137
Assert.NotNull(actual);
138138
Assert.True(actual.Count > 0);
139139
Assert.NotNull(actual.PredictiveSuggestions.First());
@@ -153,7 +153,7 @@ public void VerifyUsingCommandBasedPredictor(string userInput)
153153
public void VerifyUsingFallbackPredictor(string userInput)
154154
{
155155
var predictionContext = PredictionContext.Create(userInput);
156-
var commandAst = predictionContext.InputAst.FindAll(p => p is CommandAst, true).LastOrDefault() as CommandAst;
156+
var commandAst = predictionContext.RelatedAsts.OfType<CommandAst>().LastOrDefault();
157157
var commandName = (commandAst?.CommandElements?.FirstOrDefault() as StringConstantExpressionAst)?.Value;
158158
var inputParameterSet = new ParameterSet(commandAst);
159159
var rawUserInput = predictionContext.InputAst.Extent.Text;
@@ -166,7 +166,7 @@ public void VerifyUsingFallbackPredictor(string userInput)
166166
1,
167167
CancellationToken.None);
168168

169-
var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
169+
var actual = this._service.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
170170
Assert.NotNull(actual);
171171
Assert.True(actual.Count > 0);
172172
Assert.NotNull(actual.PredictiveSuggestions.First());
@@ -176,7 +176,7 @@ public void VerifyUsingFallbackPredictor(string userInput)
176176
Assert.Equal<string>(expected.SourceTexts, actual.SourceTexts);
177177
Assert.All<SuggestionSource>(actual.SuggestionSources, (source) => Assert.Equal(SuggestionSource.StaticCommands, source));
178178

179-
actual = this._noCommandBasedPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
179+
actual = this._noCommandBasedPredictorService.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
180180
Assert.NotNull(actual);
181181
Assert.True(actual.Count > 0);
182182
Assert.NotNull(actual.PredictiveSuggestions.First());
@@ -199,33 +199,43 @@ public void VerifyUsingFallbackPredictor(string userInput)
199199
public void VerifyNoPrediction(string userInput)
200200
{
201201
var predictionContext = PredictionContext.Create(userInput);
202-
var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
202+
var actual = this._service.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
203203
Assert.Equal(0, actual.Count);
204204

205-
actual = this._noFallbackPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
205+
actual = this._noFallbackPredictorService.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
206206
Assert.Equal(0, actual.Count);
207207

208-
actual = this._noCommandBasedPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
208+
actual = this._noCommandBasedPredictorService.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
209209
Assert.Equal(0, actual.Count);
210210

211-
actual = this._noPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
211+
actual = this._noPredictorService.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
212+
Assert.Null(actual);
213+
}
214+
215+
/// <summary>
216+
/// Verify that it returns null when we cannot parse the user input.
217+
/// </summary>
218+
[Theory]
219+
[InlineData("git status")]
220+
public void VerifyFailToParseUserInput(string userInput)
221+
{
222+
var predictionContext = PredictionContext.Create(userInput);
223+
var actual = this._service.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
212224
Assert.Null(actual);
213225
}
214226

215227
/// <summary>
216228
/// Verify when we cannot parse the user input correctly.
217229
/// </summary>
218230
/// <remarks>
219-
/// When we can parse them correctly, please move the InlineData to the corresponding test methods, for example, "git status"
220-
/// doesn't have any prediction so it should move to <see cref="VerifyNoPrediction"/>.
231+
/// When we can parse them correctly, please move the InlineData to the corresponding test methods.
221232
/// </remarks>
222233
[Theory]
223-
[InlineData("git status")]
224234
[InlineData("Get-AzContext Name")]
225235
public void VerifyMalFormattedCommandLine(string userInput)
226236
{
227237
var predictionContext = PredictionContext.Create(userInput);
228-
Action actual = () => this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
238+
Action actual = () => this._service.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
229239
_ = Assert.Throws<InvalidOperationException>(actual);
230240
}
231241
}

‎tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorTests.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ public void VerifySupportedCommandMasked()
140140
public void VerifySuggestion(string userInput)
141141
{
142142
var predictionContext = PredictionContext.Create(userInput);
143-
var expected = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
143+
var expected = this._service.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
144144
var actual = this._azPredictor.GetSuggestion(predictionContext, CancellationToken.None);
145145

146146
Assert.Equal(expected.Count, actual.Count);

‎tools/Az.Tools.Predictor/Az.Tools.Predictor.sln

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Az.Tools.Predictor", "Az.To
77
EndProject
88
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Az.Tools.Predictor.Test", "Az.Tools.Predictor.Test\Az.Tools.Predictor.Test.csproj", "{C7A3ED31-8F41-4643-ADCF-85DF032BD8AC}"
99
EndProject
10-
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "MockPSConsole", "MockPSConsole\MockPSConsole.csproj", "{80AFAFC7-9542-4CEB-AB5B-D1385A28CCE5}"
10+
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "MockPSConsole", "MockPSConsole\MockPSConsole.csproj", "{80AFAFC7-9542-4CEB-AB5B-D1385A28CCE5}"
11+
ProjectSection(ProjectDependencies) = postProject
12+
{E4A5F697-086C-4908-B90E-A31EE47ECF13} = {E4A5F697-086C-4908-B90E-A31EE47ECF13}
13+
EndProjectSection
1114
EndProject
1215
Global
1316
GlobalSection(SolutionConfigurationPlatforms) = preSolution

‎tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictor.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
using System.Runtime.CompilerServices;
2424
using System.Threading;
2525

26-
[assembly:InternalsVisibleTo("Microsoft.Azure.PowerShell.Tools.AzPredictor.Test")]
26+
[assembly: InternalsVisibleTo("Microsoft.Azure.PowerShell.Tools.AzPredictor.Test")]
2727

2828
namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
2929
{
@@ -164,7 +164,7 @@ public List<PredictiveSuggestion> GetSuggestion(PredictionContext context, Cance
164164
{
165165
var localCancellationToken = Settings.ContinueOnTimeout ? CancellationToken.None : cancellationToken;
166166

167-
suggestions = _service.GetSuggestion(context.InputAst, _settings.SuggestionCount.Value, _settings.MaxAllowedCommandDuplicate.Value, localCancellationToken);
167+
suggestions = _service.GetSuggestion(context, _settings.SuggestionCount.Value, _settings.MaxAllowedCommandDuplicate.Value, localCancellationToken);
168168

169169
var returnedValue = suggestions?.PredictiveSuggestions?.ToList();
170170
return returnedValue ?? new List<PredictiveSuggestion>();

‎tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorService.cs

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
using System.Collections.Generic;
1919
using System.Linq;
2020
using System.Management.Automation.Language;
21+
using System.Management.Automation.Subsystem;
2122
using System.Net.Http;
2223
using System.Net.Http.Headers;
2324
using System.Text;
@@ -143,22 +144,46 @@ protected virtual void Dispose(bool disposing)
143144
/// Tries to get the suggestions for the user input from the command history. If that doesn't find
144145
/// <paramref name="suggestionCount"/> suggestions, it'll fallback to find the suggestion regardless of command history.
145146
/// </remarks>
146-
public CommandLineSuggestion GetSuggestion(Ast input, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken)
147+
public CommandLineSuggestion GetSuggestion(PredictionContext context, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken)
147148
{
148-
Validation.CheckArgument(input, $"{nameof(input)} cannot be null");
149+
Validation.CheckArgument(context, $"{nameof(context)} cannot be null");
149150
Validation.CheckArgument<ArgumentOutOfRangeException>(suggestionCount > 0, $"{nameof(suggestionCount)} must be larger than 0.");
150151
Validation.CheckArgument<ArgumentOutOfRangeException>(maxAllowedCommandDuplicate > 0, $"{nameof(maxAllowedCommandDuplicate)} must be larger than 0.");
151152

152-
var commandAst = input.FindAll(p => p is CommandAst, true).LastOrDefault() as CommandAst;
153-
var commandName = (commandAst?.CommandElements?.FirstOrDefault() as StringConstantExpressionAst)?.Value;
153+
var relatedAsts = context.RelatedAsts;
154+
CommandAst commandAst = null;
155+
156+
for (var i = relatedAsts.Count - 1; i >= 0; --i)
157+
{
158+
if (relatedAsts[i] is CommandAst c)
159+
{
160+
commandAst = c;
161+
break;
162+
}
163+
}
164+
165+
var commandName = commandAst?.GetCommandName();
154166

155167
if (string.IsNullOrWhiteSpace(commandName))
156168
{
157169
return null;
158170
}
159171

160-
var inputParameterSet = new ParameterSet(commandAst);
161-
var rawUserInput = input.Extent.Text;
172+
ParameterSet inputParameterSet = null;
173+
174+
try
175+
{
176+
inputParameterSet = new ParameterSet(commandAst);
177+
}
178+
catch when (!IsSupportedCommand(commandName))
179+
{
180+
// We only ignore the exception when the command name is not supported.
181+
// For the supported ones, this most likely happens when positional parameters are used.
182+
// We want to collect the telemetry about the exception how common a positional parameter is used.
183+
return null;
184+
}
185+
186+
var rawUserInput = context.InputAst.ToString();
162187
var presentCommands = new Dictionary<string, int>();
163188
var commandBasedPredictor = _commandBasedPredictor;
164189
var commandToRequestPrediction = _commandToRequestPrediction;

‎tools/Az.Tools.Predictor/Az.Tools.Predictor/CommandLinePredictor.cs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,12 +65,13 @@ public CommandLinePredictor(IList<PredictiveCommand> modelPredictions, Parameter
6565
{
6666
var predictionText = CommandLineUtilities.EscapePredictionText(predictiveCommand.Command);
6767
Ast ast = Parser.ParseInput(predictionText, out Token[] tokens, out _);
68-
var commandAst = (ast.Find((ast) => ast is CommandAst, searchNestedScriptBlocks: false) as CommandAst);
68+
var commandAst = ast.Find((ast) => ast is CommandAst, searchNestedScriptBlocks: false) as CommandAst;
69+
var commandName = commandAst?.GetCommandName();
6970

70-
if (commandAst?.CommandElements[0] is StringConstantExpressionAst commandName)
71+
if (!string.IsNullOrWhiteSpace(commandName))
7172
{
7273
var parameterSet = new ParameterSet(commandAst);
73-
this._commandLinePredictions.Add(new CommandLine(commandName.Value, predictiveCommand.Description, parameterSet));
74+
this._commandLinePredictions.Add(new CommandLine(commandName, predictiveCommand.Description, parameterSet));
7475
}
7576
}
7677
}

‎tools/Az.Tools.Predictor/Az.Tools.Predictor/IAzPredictorService.cs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212
// limitations under the License.
1313
// ----------------------------------------------------------------------------------
1414

15-
using System;
1615
using System.Collections.Generic;
1716
using System.Management.Automation.Language;
17+
using System.Management.Automation.Subsystem;
1818
using System.Threading;
1919

2020
namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
@@ -27,12 +27,12 @@ public interface IAzPredictorService
2727
/// <summary>
2828
/// Gest the suggestions for the user input.
2929
/// </summary>
30-
/// <param name="input">User input from PSReadLine.</param>
30+
/// <param name="context">User input context from PSReadLine.</param>
3131
/// <param name="suggestionCount">The number of suggestion to return.</param>
3232
/// <param name="cancellationToken">The cancellation token</param>
3333
/// <param name="maxAllowedCommandDuplicate">The maximum amount of the same commnds in the list of predictions.</param>
34-
/// <returns>The suggestions for <paramref name="input"/>. The maximum number of suggestions is <paramref name="suggestionCount"/>.</returns>
35-
public CommandLineSuggestion GetSuggestion(Ast input, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken);
34+
/// <returns>The suggestions for <paramref name="context"/>. The maximum number of suggestions is <paramref name="suggestionCount"/>.</returns>
35+
public CommandLineSuggestion GetSuggestion(PredictionContext context, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken);
3636

3737
/// <summary>
3838
/// Requests predictions, given a command string.

‎tools/Az.Tools.Predictor/Az.Tools.Predictor/ParameterSet.cs

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
2424
/// does not matter to resulting prediction - the prediction should adapt to the
2525
/// order of the parameters typed by the user.
2626
/// </summary>
27+
/// <remarks>
28+
/// This doesn't handle the positional parameters yet.
29+
/// </remarks>
2730
sealed class ParameterSet
2831
{
2932
/// <summary>
@@ -36,11 +39,14 @@ public ParameterSet(CommandAst commandAst)
3639
Validation.CheckArgument(commandAst, $"{nameof(commandAst)} cannot be null.");
3740

3841
var parameters = new List<Parameter>();
39-
var elements = commandAst.CommandElements.Skip(1);
4042
CommandParameterAst param = null;
4143
Ast arg = null;
42-
foreach (Ast elem in elements)
44+
45+
// Loop through all the parameters. The first element of CommandElements is the command name, so skip it.
46+
for (var i = 1; i < commandAst.CommandElements.Count(); ++i)
4347
{
48+
var elem = commandAst.CommandElements[i];
49+
4450
if (elem is CommandParameterAst p)
4551
{
4652
AddParameter(param, arg);
@@ -68,11 +74,22 @@ public ParameterSet(CommandAst commandAst)
6874

6975
Parameters = parameters;
7076

71-
void AddParameter(CommandParameterAst parameterName, Ast parameterValue)
77+
void AddParameter(CommandParameterAst parameter, Ast parameterValue)
7278
{
73-
if (parameterName != null)
79+
if (parameter != null)
7480
{
75-
parameters.Add(new Parameter(parameterName.ParameterName, (parameterValue == null) ? null : CommandLineUtilities.UnescapePredictionText(parameterValue.ToString())));
81+
var value = parameterValue?.ToString();
82+
if (value == null)
83+
{
84+
value = parameter.Argument?.ToString();
85+
}
86+
87+
if (value != null)
88+
{
89+
value = CommandLineUtilities.UnescapePredictionText(value);
90+
}
91+
92+
parameters.Add(new Parameter(parameter.ParameterName, value));
7693
}
7794
}
7895
}

0 commit comments

Comments
 (0)