LLC

autorest/codegen/__init__.py

@@ -116,7 +116,6 @@ def _create_code_model(self, yaml_data: Dict[str, Any], options: Dict[str, Union
             code_model.generate_single_parameter_from_multiple_media_types()
 
         # LRO operation
-        code_model.format_lro_operations()
         code_model.remove_next_operation()
 
         if options["credential"]:

autorest/codegen/models/__init__.py

@@ -23,6 +23,7 @@
 from .operation_group import OperationGroup
 from .schema_response import SchemaResponse
 from .parameter_list import ParameterList
+from .request import Request
 
 
 __all__ = [
@@ -46,6 +47,7 @@
     "ParameterList",
     "OperationGroup",
     "Property",
+    "Request",
     "SchemaResponse",
     "TokenCredentialSchema",
 ]

autorest/codegen/models/code_model.py

@@ -143,40 +143,6 @@ def add_credential_global_parameter(self) -> None:
         )
         self.global_parameters.insert(0, credential_parameter)
 
-    @staticmethod
-    def _lro_initial_function(operation: LROOperation) -> Operation:
-        return Operation(
-            yaml_data={},
-            name="_" + operation.name + "_initial",
-            description="",
-            url=operation.url,
-            method=operation.method,
-            multipart=operation.multipart,
-            api_versions=operation.api_versions,
-            parameters=operation.parameters.parameters,
-            requests=operation.requests,
-            responses=operation.responses,
-            exceptions=operation.exceptions,
-            want_description_docstring=False,
-            want_tracing=False,
-        )
-
-    def format_lro_operations(self) -> None:
-        """Adds operations and attributes needed for LROs.
-
-        If there are LRO functions in here, will add initial LRO function. Will also set the return
-        type of the LRO operation
-        """
-        for operation_group in self.operation_groups:
-            i = 0
-            while i < len(operation_group.operations):
-                operation = operation_group.operations[i]
-                if isinstance(operation, LROOperation):
-                    operation.set_lro_response_type()
-                    operation_group.operations.insert(i, CodeModel._lro_initial_function(operation))
-                    i += 1
-                i += 1
-
     def remove_next_operation(self) -> None:
         """Linking paging operations together.
         """
@@ -302,11 +268,12 @@ def add_schema_link_to_operation(self) -> None:
             for operation in operation_group.operations:
                 for obj in chain(
                     operation.parameters,
-                    operation.multiple_media_type_parameters or [],
+                    operation.request.parameters,
+                    operation.request.multiple_media_type_parameters or [],
                     operation.responses,
                     operation.exceptions,
                     chain.from_iterable(response.headers for response in operation.responses),
-                    chain.from_iterable(request.parameters for request in operation.requests)
+                    chain.from_iterable(request.parameters for request in operation.request.schema_requests)
                 ):
                     self._populate_schema(obj)
 
@@ -317,15 +284,16 @@ def add_schema_link_to_global_parameters(self) -> None:
     def generate_single_parameter_from_multiple_media_types(self) -> None:
         for operation_group in self.operation_groups:
             for operation in operation_group.operations:
-                if operation.multiple_media_type_parameters:
+                if operation.request.multiple_media_type_parameters:
                     type_annot = ", ".join([
-                        param.schema.operation_type_annotation for param in operation.multiple_media_type_parameters
+                        param.schema.operation_type_annotation
+                        for param in operation.request.multiple_media_type_parameters
                     ])
                     docstring_type = " or ".join([
-                        param.schema.docstring_type for param in operation.multiple_media_type_parameters
+                        param.schema.docstring_type for param in operation.request.multiple_media_type_parameters
                     ])
                     chosen_parameter = next(
-                        iter(filter(lambda x: x.has_multiple_media_types, operation.parameters)), None
+                        iter(filter(lambda x: x.has_multiple_media_types, operation.request.parameters)), None
                     )
                     if not chosen_parameter:
                         raise ValueError("You are missing a parameter that has multiple media types")

autorest/codegen/models/lro_operation.py

@@ -9,7 +9,7 @@
 from .operation import Operation
 from .parameter import Parameter
 from .schema_response import SchemaResponse
-from .schema_request import SchemaRequest
+from .request import Request
 from .imports import ImportType, TypingSection
 from .base_schema import BaseSchema
 
@@ -20,44 +20,39 @@ class LROOperation(Operation):
     def __init__(
         self,
         yaml_data: Dict[str, Any],
+        request: Request,
         name: str,
         description: str,
-        url: str,
-        method: str,
-        multipart: bool,
         api_versions: Set[str],
-        requests: List[SchemaRequest],
-        summary: Optional[str] = None,
         parameters: Optional[List[Parameter]] = None,
         multiple_media_type_parameters: Optional[List[Parameter]] = None,
+        summary: Optional[str] = None,
         responses: Optional[List[SchemaResponse]] = None,
         exceptions: Optional[List[SchemaResponse]] = None,
         want_description_docstring: bool = True,
         want_tracing: bool = True
     ) -> None:
         super(LROOperation, self).__init__(
             yaml_data,
+            request,
             name,
             description,
-            url,
-            method,
-            multipart,
             api_versions,
-            requests,
-            summary,
             parameters,
             multiple_media_type_parameters,
+            summary,
             responses,
             exceptions,
             want_description_docstring,
             want_tracing,
         )
-        self.lro_response: Optional[SchemaResponse] = None
         self.lro_options = yaml_data.get("extensions", {}).get("x-ms-long-running-operation-options", {})
+        self.name = "begin_" + self.name
 
-    def set_lro_response_type(self) -> None:
+    @property
+    def lro_response(self) -> Optional[SchemaResponse]:
         if not self.responses:
-            return
+            return None
         responses_with_bodies = [r for r in self.responses if r.has_body]
         num_response_schemas = {r.schema for r in responses_with_bodies}
         response = None
@@ -81,7 +76,7 @@ def set_lro_response_type(self) -> None:
 
         elif num_response_schemas:
             response = responses_with_bodies[0]
-        self.lro_response = response
+        return response
 
     @property
     def has_optional_return_type(self) -> bool:

autorest/codegen/models/lro_paging_operation.py

@@ -3,52 +3,11 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-
-from typing import Any, Dict, List, Set, Optional
+from .imports import FileImport
 from .lro_operation import LROOperation
 from .paging_operation import PagingOperation
-from .imports import FileImport
-from .schema_request import SchemaRequest
-from .parameter import Parameter
-from .schema_response import SchemaResponse
 
 class LROPagingOperation(PagingOperation, LROOperation):
-    def __init__(
-        self,
-        yaml_data: Dict[str, Any],
-        name: str,
-        description: str,
-        url: str,
-        method: str,
-        multipart: bool,
-        api_versions: Set[str],
-        requests: List[SchemaRequest],
-        summary: Optional[str] = None,
-        parameters: Optional[List[Parameter]] = None,
-        multiple_media_type_parameters: Optional[List[Parameter]] = None,
-        responses: Optional[List[SchemaResponse]] = None,
-        exceptions: Optional[List[SchemaResponse]] = None,
-        want_description_docstring: bool = True,
-        want_tracing: bool = True
-    ) -> None:
-        super(LROPagingOperation, self).__init__(
-            yaml_data,
-            name,
-            description,
-            url,
-            method,
-            multipart,
-            api_versions,
-            requests,
-            summary,
-            parameters,
-            multiple_media_type_parameters,
-            responses,
-            exceptions,
-            want_description_docstring,
-            want_tracing,
-            override_success_response_to_200=True
-        )
 
     def imports(self, code_model, async_mode: bool) -> FileImport:
         lro_imports = LROOperation.imports(self, code_model, async_mode)
@@ -57,3 +16,8 @@ def imports(self, code_model, async_mode: bool) -> FileImport:
         file_import = lro_imports
         file_import.merge(paging_imports)
         return file_import
+
+    @property
+    def paging_success_status_code(self):
+        # hardcode paging's success status code in lro + paging to be 200
+        return [200]