fix bugs

Author: Glauco Knihs

RoyalCode.Utils/RoyalCode.Extensions.SourceGenerator/Generators/ArgumentsGenerator.cs

@@ -2,7 +2,7 @@
 
 namespace RoyalCode.Extensions.SourceGenerator.Generators;
 
-public class ArgumentsGenerator : GeneratorNode
+public class ArgumentsGenerator : GeneratorNode, IWithNamespaces
 {
     private List<ValueNode>? arguments;
 
@@ -20,6 +20,16 @@ public void AddArguments(ValueNode[] arguments)
         this.arguments.AddRange(arguments);
     }
 
+    public IEnumerable<string> GetNamespaces()
+    {
+        if (arguments is null)
+            yield break;
+        foreach (var arg in arguments)
+            if (arg is IWithNamespaces withNamespaces)
+                foreach (var ns in withNamespaces.GetNamespaces())
+                    yield return ns;
+    }
+
     public override void Write(StringBuilder sb, int ident = 0)
     {
         sb.Append("(");

RoyalCode.Utils/RoyalCode.Extensions.SourceGenerator/Generators/GeneratorNodeList.cs

@@ -18,9 +18,12 @@ public void Add(GeneratorNode generator)
 
     public IEnumerable<string> GetNamespaces()
     {
-        if (nodes is not null)
-            foreach (var ns in nodes.OfType<IWithNamespaces>().SelectMany(node => node.GetNamespaces()))
-                yield return ns;
+        if (nodes is null)
+            yield break;
+        foreach (var node in nodes)
+            if (node is IWithNamespaces withNamespaces)
+                foreach (var ns in withNamespaces.GetNamespaces())
+                    yield return ns;
     }
 
     public override void Write(StringBuilder sb, int ident = 0)

RoyalCode.Utils/RoyalCode.Extensions.SourceGenerator/Generators/MethodInvokeGenerator.cs

@@ -2,7 +2,7 @@
 
 namespace RoyalCode.Extensions.SourceGenerator.Generators;
 
-public class MethodInvokeGenerator : GeneratorNode
+public class MethodInvokeGenerator : GeneratorNode, IWithNamespaces
 {
     private readonly ValueNode identifier;
     private readonly string method;
@@ -30,6 +30,15 @@ public MethodInvokeGenerator(ValueNode identifier, string method, ValueNode argu
 
     public void AddArgument(ValueNode arg) => arguments.AddArgument(arg);
 
+    public IEnumerable<string> GetNamespaces()
+    {
+        if (identifier is IWithNamespaces withNamespaces)
+            foreach (var ns in withNamespaces.GetNamespaces())
+                yield return ns;
+        foreach (var ns in arguments.GetNamespaces())
+            yield return ns;
+    }
+
     public void UseArgumentPerLine() => arguments.InLine = false;
 
     public override void Write(StringBuilder sb, int ident = 0)