# -------------------------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for license information. # -------------------------------------------------------------------------------------------- from enum import Enum import importlib import inspect import logging from pathlib import Path import pkgutil import shutil import sys import tempfile from msrest.serialization import Model from msrest.paging import Paged _LOGGER = logging.getLogger(__name__) copyright_header = b"""# coding=utf-8 # -------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. # # Code generated by Microsoft (R) AutoRest Code Generator. # Changes may cause incorrect behavior and will be lost if the code is # regenerated. # -------------------------------------------------------------------------- """ header = copyright_header + b"""from msrest.serialization import Model from msrest.exceptions import HttpOperationError """ track2_header = copyright_header + b"""import datetime import msrest import msrest.serialization from typing import Dict, List, Optional, Union from msrest.exceptions import HttpOperationError """ paging_header = copyright_header + b"""from msrest.paging import Paged """ init_file = """ try: from .{} import * except (SyntaxError, ImportError): from .{} import * from .{} import * """ def parse_input(input_parameter): """From a syntax like package_name#submodule, build a package name and complete module name. """ split_package_name = input_parameter.split('#') package_name = split_package_name[0] module_name = package_name.replace("-", ".") if len(split_package_name) >= 2: module_name = ".".join([module_name, split_package_name[1]]) return package_name, module_name def solve_mro(models, track2=False): for models_module in models: models_path = models_module.__path__[0] _LOGGER.info("Working on %s", models_path) if not track2 and Path(models_path, "models_py3.py").exists(): _LOGGER.info("Skipping since already patched") return # Build the new files in a temp folder with tempfile.TemporaryDirectory() as temp_folder: final_models_path = Path(temp_folder, "models") final_models_path.mkdir() solve_one_model(models_module, final_models_path, track2=track2) # Switch the files shutil.rmtree(models_path) shutil.move(final_models_path, models_path) def solve_one_model(models_module, output_folder, track2=False): """Will build the compacted models in the output_folder""" models_classes = [ (len(model_class.__mro__), inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items() if model_name[0].isupper() and Model in model_class.__mro__ ] # Only sort based on the first element in the tuple models_classes.sort(key=lambda x: x[0]) paged_models_classes = [ (inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items() if model_name[0].isupper() and Paged in model_class.__mro__ ] enum_models_classes = [ (inspect.getfile(model_class), model_class) for model_name, model_class in vars(models_module).items() if model_name[0].isupper() and Enum in model_class.__mro__ ] if enum_models_classes: enum_file = Path(enum_models_classes[0][0]) shutil.copyfile(enum_file, Path(output_folder, enum_file.name)) enum_file_module_name = enum_file.with_suffix('').name else: enum_file_module_name = None write_model_file(Path(output_folder, "models_py3.py"), models_classes, track2=track2) write_paging_file(Path(output_folder, "paged_models.py"), paged_models_classes) write_init( Path(output_folder, "__init__.py"), "models_py3", "models", "paged_models", enum_file_module_name ) def write_model_file(output_file_path, classes_to_write, track2=False): with open(output_file_path, "bw") as write_fd: if track2: write_fd.write(track2_header) else: write_fd.write(header) for model in classes_to_write: _, model_file_path, _ = model with open(model_file_path, "rb") as read_fd: lines = read_fd.readlines() # Skip until it's "class XXXX" while lines: if lines[0].startswith(b"class "): break lines.pop(0) else: raise ValueError("Never found any class definition!") # Now I keep everything write_fd.write(b'\n') write_fd.write(b'\n') write_fd.writelines(lines) def write_paging_file(output_file_path, classes_to_write): with open(output_file_path, "bw") as write_fd: write_fd.write(paging_header) for model in classes_to_write: model_file_path, _ = model with open(model_file_path, "rb") as read_fd: # Skip the first 15 lines (based on Autorest deterministic behavior) # If we want this less random, look for the first line starts with "class" lines = read_fd.readlines()[14:] write_fd.write(b'\n') write_fd.write(b'\n') write_fd.writelines(lines) def write_init(output_file_path, model_file_name, model_file_name_py2, paging_file_name, enum_file_name): with open(output_file_path, "bw") as write_fd: write_fd.write(copyright_header) write_fd.write(init_file.format( model_file_name, model_file_name_py2, paging_file_name, ).encode('utf8')) if enum_file_name: write_fd.write( "from .{} import *".format(enum_file_name).encode('utf8') ) def find_models_to_change(module_name): """Will figure out if the package is a multi-api one, and understand what to generate. """ main_module = importlib.import_module(module_name) try: models_module = main_module.models models_module.__path__ # It didn't fail, that's a single API package return [models_module] except AttributeError: # This means I loaded the fake module "models" # and it's multi-api, load all models return [ importlib.import_module('.' + label + '.models', main_module.__name__) for (_, label, ispkg) in pkgutil.iter_modules(main_module.__path__) if ispkg and label != 'aio' ] def find_autorest_generated_folder(module_prefix="azure.mgmt"): """Find all Autorest generated code in that module prefix. This actually looks for a "models" package only. We could be smarter if necessary. """ _LOGGER.info("Looking for Autorest generated package in %s", module_prefix) result = [] prefix_module = importlib.import_module(module_prefix) for _, sub_package, ispkg in pkgutil.iter_modules(prefix_module.__path__, module_prefix + "."): try: _LOGGER.debug("Try %s", sub_package) importlib.import_module(".models", sub_package) # If not exception, we found it _LOGGER.info("Found %s", sub_package) result.append(sub_package) except ModuleNotFoundError: # No model, might dig deeper if ispkg: result += find_autorest_generated_folder(sub_package) return result if __name__ == "__main__": logging.basicConfig(level=logging.INFO) track2_packages = [ 'azure.mgmt.keyvault', 'azure.mgmt.storage', 'azure.mgmt.compute', 'azure.mgmt.monitor', 'azure.mgmt.rdbms' 'azure.mgmt.loganalytics', 'azure.mgmt.web', 'azure.mgmt.cosmosdb', 'azure.mgmt.privatedns', 'azure.mgmt.dms', 'azure.mgmt.sqlvirtualmachine' ] prefix = sys.argv[1] if len(sys.argv) >= 2 else "azure.mgmt" for autorest_package in find_autorest_generated_folder(prefix): models = find_models_to_change(autorest_package) track2 = False for track2_pkg in track2_packages: if autorest_package.startswith(track2_pkg): track2 = True break solve_mro(models, track2=track2)