Add -CountVariable parameter. (#352)

Author: peombwa

src/readme.graph.md

@@ -469,19 +469,27 @@ directive:
       {
         return $;
       } else {
-        // Add custom -PageSize parameter to *_List cmdlets that support Odata next link.
         let odataNextLinkRegex = /(^\s*)(if\s*\(\s*result.OdataNextLink\s*!=\s*null\s*\))/gmi
         if($.match(odataNextLinkRegex)) {
+          // Add custom -PageSize parameter to *_List cmdlets that support Odata next link.
           $ = $.replace(odataNextLinkRegex, '$1if (result.OdataNextLink != null && this.ShouldIteratePages(this.InvocationInformation.BoundParameters, result.Value.Length))\n$1');
 
           let psBaseClassImplementationRegex = /(\s*:\s*)(global::System.Management.Automation.PSCmdlet)/gmi
           $ = $.replace(psBaseClassImplementationRegex, '$1Microsoft.Graph.PowerShell.Cmdlets.Custom.ListCmdlet');
 
           let beginProcessingRegex = /(^\s*)(protected\s*override\s*void\s*BeginProcessing\(\)\s*{)/gmi
-          $ = $.replace(beginProcessingRegex, '$1$2\n$1  if (this.InvocationInformation?.BoundParameters != null){ InitializePaging(ref this.__invocationInfo, ref this._top); }\n$1');
+          $ = $.replace(beginProcessingRegex, '$1$2\n$1  if (this.InvocationInformation?.BoundParameters != null){ InitializeCmdlet(ref this.__invocationInfo, ref this._top, ref this._count); }\n$1');
 
           let odataNextLinkCallRegex = /(^\s*)(await\s*this\.Client\.UsersUserListUser_Call\(requestMessage\,\s*onOk\,\s*onDefault\,\s*this\,\s*Pipeline\)\;)/gmi
           $ = $.replace(odataNextLinkCallRegex, '$1requestMessage.RequestUri = GetOverflowItemsNextLinkUri(requestMessage.RequestUri);\n$1$2');
+
+          // Set -Count parameter to private. This will be replaced by -CountVariable
+          let countParameterRegex = /public(\s*global::System\.Management\.Automation\.SwitchParameter\s*Count\s*)/gm
+          $ = $.replace(countParameterRegex, 'private$1');
+
+          // Call OnBeforeWriteObject to deserialize '@odata.count' from the response object.
+          let writeObjectRegex = /^(\s*)(WriteObject\(result\.Value,true\);)$/gm
+          $ = $.replace(writeObjectRegex,'\n$1OnBeforeWriteObject(this.InvocationInformation.BoundParameters, result?.AdditionalProperties);\n$1$2');
         }
         return $;
       }
@@ -502,4 +510,18 @@ directive:
         $ = $.replace(propertyContainsRegex, '$1$2, System.StringComparer.OrdinalIgnoreCase$3');
         return $;
       }
+
+# Serialize all $count parameter to lowercase true or false.
+  - from: source-file-csharp
+    where: $
+    transform: >
+      if (!$documentPath.match(/generated%5Capi%5C\w*.cs/gm))
+      {
+        return $;
+      } else {
+        // Add '.ToLower()' at the end of all 'Count.ToString()'
+        let countRegex = /(Count\.ToString\(\))/gmi
+        $ = $.replace(countRegex, '$1.ToLower()');
+        return $;
+      }
 ```

tools/Custom/ListCmdlet.cs

@@ -34,6 +34,20 @@ public partial class ListCmdlet : global::System.Management.Automation.PSCmdlet
         [global::Microsoft.Graph.PowerShell.Category(global::Microsoft.Graph.PowerShell.ParameterCategory.Runtime)]
         public global::System.Management.Automation.SwitchParameter All { get => this._all; set => this._all = value; }
 
+        // <summary>Backing field for <see cref="CountVariable" /> property.</summary>
+        private string _countVariable;
+
+        /// <summary>Specifies a count of the total number of items in a collection. </summary>
+        [global::System.Management.Automation.Parameter(Mandatory = false, HelpMessage = "Specifies a count of the total number of items in a collection. By default, this variable will be set in the global scope.")]
+        [Microsoft.Graph.PowerShell.Runtime.Info(
+        Required = false,
+        ReadOnly = false,
+        Description = @"Specifies a count of the total number of items in a collection. By default, this variable will be set in the global scope.",
+        PossibleTypes = new[] { typeof(string) })]
+        [global::Microsoft.Graph.PowerShell.Category(global::Microsoft.Graph.PowerShell.ParameterCategory.Runtime)]
+        [global::System.Management.Automation.Alias("CV")]
+        public string CountVariable { get => this._countVariable; set => this._countVariable = value; }
+
         /// <summary>
         /// Default number of items per page.
         /// </summary>
@@ -75,7 +89,7 @@ public partial class ListCmdlet : global::System.Management.Automation.PSCmdlet
         /// </summary>
         /// <param name="invocationInfo">A reference to <see cref="System.Management.Automation.InvocationInfo"/> object.</param>
         /// <param name="top">A reference to top parameter.</param>
-        public void InitializePaging(ref global::System.Management.Automation.InvocationInfo invocationInfo, ref int top)
+        public void InitializeCmdlet(ref global::System.Management.Automation.InvocationInfo invocationInfo, ref int top, ref global::System.Management.Automation.SwitchParameter count)
         {
             if (invocationInfo.BoundParameters.ContainsKey("PageSize") && (PageSize > MaxPageSize || PageSize == default))
             {
@@ -110,6 +124,13 @@ public void InitializePaging(ref global::System.Management.Automation.Invocation
                 requiredPages = limit / currentPageSize;
                 overflowItemsCount = limit % currentPageSize;
             }
+
+            if ((!invocationInfo.BoundParameters.ContainsKey("Count")) && invocationInfo.BoundParameters.ContainsKey("CountVariable"))
+            {
+                // Set Count to true when CountVariable is set.
+                invocationInfo.BoundParameters["Count"] = true;
+                count = true;
+            }
         }
 
         /// <summary>
@@ -150,5 +171,20 @@ public bool ShouldIteratePages(global::System.Collections.Generic.Dictionary<str
             }
             return nextLinkUri.Uri;
         }
+
+        internal void OnBeforeWriteObject(global::System.Collections.Generic.Dictionary<string, object> boundParameters, global::System.Collections.Generic.IDictionary<string, object> additionalProperties)
+        {
+            // Get odata.count from the response.
+            if (boundParameters.ContainsKey("CountVariable") &&
+                additionalProperties != null &&
+                additionalProperties.TryGetValue("@odata.count", out var odataCount))
+            {
+                // Save the Count back to the PS environment in a global variable.
+                // We need to store count in a global variable since these cmdlets are exported as functions.
+                // i.e. Functions can't modify parent scope.
+                var psVI = SessionState.PSVariable;
+                psVI.Set(new PSVariable(CountVariable.Contains(":") ? CountVariable : $"global:{CountVariable}", odataCount));
+            }
+        }
     }
 }