Allow mutation of InvokeMethodExpression and TryCatchFinallyStatement (#7229)

JoshLove-msft · web-flow · commit cb5f88de81d7 · 2025-05-06T18:41:21.000Z
Contributes to Azure/azure-sdk-for-net#49632
diff --git a/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/Expressions/InvokeMethodExpression.cs b/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/Expressions/InvokeMethodExpression.cs
@@ -21,31 +21,71 @@ private InvokeMethodExpression(
             bool addConfigureAwaitFalse = true,
             CSharpType? extensionType = null)
         {
-            InstanceReference = instanceReference;
-            MethodName = methodName;
-            MethodSignature = methodSignature;
-            Arguments = arguments;
-            TypeArguments = typeArguments;
-            CallAsAsync = callAsAsync;
-            AddConfigureAwaitFalse = addConfigureAwaitFalse;
-            ExtensionType = extensionType;
+            _instanceReference = instanceReference;
+            _methodName = methodName;
+            _methodSignature = methodSignature;
+            _arguments = arguments;
+            _typeArguments = typeArguments;
+            _callAsAsync = callAsAsync;
+            _addConfigureAwaitFalse = addConfigureAwaitFalse;
+            _extensionType = extensionType;
         }
 
-        public CSharpType? ExtensionType { get; init; }
+        public CSharpType? ExtensionType
+        {
+            get => _extensionType;
+            init => _extensionType = value;
+        }
+        private CSharpType? _extensionType;
 
-        public bool AddConfigureAwaitFalse { get; init; }
+        public bool AddConfigureAwaitFalse
+        {
+            get => _addConfigureAwaitFalse;
+            init => _addConfigureAwaitFalse = value;
+        }
+        private bool _addConfigureAwaitFalse;
 
-        public bool CallAsAsync { get; init; }
+        public bool CallAsAsync
+        {
+            get => _callAsAsync;
+            init => _callAsAsync = value;
+        }
+        private bool _callAsAsync;
 
-        public IReadOnlyList<CSharpType>? TypeArguments { get; init; }
+        public IReadOnlyList<CSharpType>? TypeArguments
+        {
+            get => _typeArguments;
+            init => _typeArguments = value;
+        }
+        private IReadOnlyList<CSharpType>? _typeArguments;
 
-        public IReadOnlyList<ValueExpression> Arguments { get; init; }
+        public IReadOnlyList<ValueExpression> Arguments
+        {
+            get => _arguments;
+            init => _arguments = value;
+        }
+        private IReadOnlyList<ValueExpression> _arguments;
 
-        public MethodSignatureBase? MethodSignature { get; init; }
+        public MethodSignatureBase? MethodSignature
+        {
+            get => _methodSignature;
+            init => _methodSignature = value;
+        }
+        private MethodSignatureBase? _methodSignature;
 
-        public string? MethodName { get; init; }
+        public string? MethodName
+        {
+            get => _methodName;
+            init => _methodName = value;
+        }
+        private string? _methodName;
 
-        public ValueExpression? InstanceReference { get; init; }
+        public ValueExpression? InstanceReference
+        {
+            get => _instanceReference;
+            init => _instanceReference = value;
+        }
+        private ValueExpression? _instanceReference;
 
         public InvokeMethodExpression(ValueExpression? instanceReference, string methodName, IReadOnlyList<ValueExpression> arguments) : this(instanceReference, methodName, null, arguments, null, false) { }
 
@@ -71,5 +111,49 @@ internal override void Write(CodeWriter writer)
 
         private MethodBodyStatement? _terminated;
         public MethodBodyStatement Terminate() => _terminated ??= new ExpressionStatement(this);
+
+        public void Update(
+            ValueExpression? instanceReference = null,
+            string? methodName = null,
+            MethodSignatureBase? methodSignature = null,
+            IReadOnlyList<ValueExpression>? arguments = null,
+            IReadOnlyList<CSharpType>? typeArguments = null,
+            bool? callAsAsync = null,
+            bool? addConfigureAwaitFalse = null,
+            CSharpType? extensionType = null)
+        {
+            if (instanceReference != null)
+            {
+                _instanceReference = instanceReference;
+            }
+            if (methodName != null)
+            {
+                _methodName = methodName;
+            }
+            if (methodSignature != null)
+            {
+                _methodSignature = methodSignature;
+            }
+            if (arguments != null)
+            {
+                _arguments = arguments;
+            }
+            if (typeArguments != null)
+            {
+                _typeArguments = typeArguments;
+            }
+            if (callAsAsync != null)
+            {
+                _callAsAsync = callAsAsync.Value;
+            }
+            if (addConfigureAwaitFalse != null)
+            {
+                _addConfigureAwaitFalse = addConfigureAwaitFalse.Value;
+            }
+            if (extensionType != null)
+            {
+                _extensionType = extensionType;
+            }
+        }
     }
 }
diff --git a/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/Statements/TryCatchFinallyStatement.cs b/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/src/Statements/TryCatchFinallyStatement.cs
@@ -9,9 +9,9 @@ namespace Microsoft.TypeSpec.Generator.Statements
 {
     public sealed class TryCatchFinallyStatement : MethodBodyStatement
     {
-        public TryStatement Try { get; }
-        public IReadOnlyList<CatchStatement> Catches { get; }
-        public FinallyStatement? Finally { get; }
+        public TryStatement Try { get; private set; }
+        public IReadOnlyList<CatchStatement> Catches { get; private set; }
+        public FinallyStatement? Finally { get; private set; }
 
         public TryCatchFinallyStatement(TryStatement @try, IReadOnlyList<CatchStatement> catches, FinallyStatement? @finally)
         {
@@ -39,5 +39,24 @@ internal override void Write(CodeWriter writer)
 
             Finally?.Write(writer);
         }
+
+        public void Update(
+            TryStatement? @try = null,
+            IReadOnlyList<CatchStatement>? catches = null,
+            FinallyStatement? @finally = null)
+        {
+            if (@try != null)
+            {
+                Try = @try;
+            }
+            if (catches != null)
+            {
+                Catches = catches;
+            }
+            if (@finally != null)
+            {
+                Finally = @finally;
+            }
+        }
     }
 }
diff --git a/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/Expressions/ExpressionsTests.cs b/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/Expressions/ExpressionsTests.cs
@@ -58,6 +58,44 @@ public void TestInvokeInstanceMethodExpression()
             Assert.AreEqual(expectedResult, test);
         }
 
+        [Test]
+        public void TestUpdateInvokeInstanceMethodExpression()
+        {
+            // declare the instance method
+            var mockTypeProvider = new Mock<TypeProvider>();
+            var returnInstanceMethod = Return(new InvokeMethodExpression(null, "Bar", [Bool(true), Bool(false), Bool(false)]));
+            var fooMethod = new MethodProvider(
+                new MethodSignature(
+                    Name: "Foo",
+                    Modifiers: MethodSignatureModifiers.Public,
+                    ReturnType: typeof(bool),
+                    Parameters: [],
+                    Description: null, ReturnDescription: null),
+                new MethodBodyStatement[] { returnInstanceMethod },
+                mockTypeProvider.Object);
+
+            var returnExpression = returnInstanceMethod as ExpressionStatement;
+            Assert.IsNotNull(returnExpression);
+
+            var keyWordExpression = returnExpression!.Expression as KeywordExpression;
+            Assert.IsNotNull(keyWordExpression);
+
+            var invokeMethodExpression = keyWordExpression!.Expression as InvokeMethodExpression;
+            Assert.IsNotNull(invokeMethodExpression);
+
+            invokeMethodExpression!.Update(
+                methodName: "Baz",
+                arguments: [Bool(true)]);
+
+            // Verify the expected behavior
+            using var writer = new CodeWriter();
+            writer.WriteMethod(fooMethod);
+
+            var expectedResult = Helpers.GetExpectedFromFile();
+            var test = writer.ToString(false);
+            Assert.AreEqual(expectedResult, test);
+        }
+
         [Test]
         public void TestNewInstanceExpression()
         {
diff --git a/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/Expressions/TestData/ExpressionsTests/TestUpdateInvokeInstanceMethodExpression.cs b/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/Expressions/TestData/ExpressionsTests/TestUpdateInvokeInstanceMethodExpression.cs
@@ -0,0 +1,4 @@
+public bool Foo()
+{
+    return Baz(true);
+}
diff --git a/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/Statements/StatementTests.cs b/packages/http-client-csharp/generator/Microsoft.TypeSpec.Generator/test/Statements/StatementTests.cs
@@ -457,6 +457,26 @@ public void TryCatchFinallyStatementWithMultipleCatches()
             Assert.AreEqual(expectedResult, test);
         }
 
+        [Test]
+        public void TryCatchFinallyUpdate()
+        {
+            var tryCatchFinally = new TryCatchFinallyStatement(new TryStatement(), new CatchStatement(null), new FinallyStatement());
+
+            var tryStatement = new TryStatement();
+            var catchStatement = new CatchStatement(null);
+            var finallyStatement = new FinallyStatement();
+
+            tryCatchFinally.Update(
+                @try: tryStatement,
+                catches: [catchStatement],
+                @finally: finallyStatement);
+
+            Assert.AreEqual(tryStatement, tryCatchFinally.Try);
+            Assert.AreEqual(1, tryCatchFinally.Catches.Count);
+            Assert.AreEqual(catchStatement, tryCatchFinally.Catches[0]);
+            Assert.AreEqual(finallyStatement, tryCatchFinally.Finally);
+        }
+
         [Test]
         public void TestIfElsePreprocessorStatement()
         {

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +public bool Foo()
 +{
 +    return Baz(true);
 +}