lmazuel/patch_models.py

## patch_models.py
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
"""

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):
    for models_module in models:
        models_path = models_module.__path__[0]
        _LOGGER.info(f"Working on {models_path}")
        if 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)

            # Swith the files
            shutil.rmtree(models_path)
            shutil.move(final_models_path, models_path)

def solve_one_model(models_module, output_folder):
    """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])

    py2_models_classes = [
        (len_mro, path.replace("_py3.py", ".py"), None)
        for len_mro, path, _ in models_classes
    ]

    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)
    write_model_file(Path(output_folder, "models.py"), py2_models_classes)
    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):
    with open(output_file_path, "bw") as write_fd:
        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
        ]

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(f"Looking for Autorest generated package in {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(f"Try {sub_package}")
            importlib.import_module(".models", sub_package)
            # If not exception, we found it
            _LOGGER.info(f"Found {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)

    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)
        solve_mro(models)

## patch_models_v2.py
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.exceptions import HttpOperationError
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
"""

paging_header = copyright_header+b"""from msrest.paging import Paged
"""

init_file = """
try:
    from .{} import *
except (SyntaxError, ImportError):
    from .{} import *
from .{} import *
"""

MODEL_PY2_NAME = "_models"
MODEL_PY3_NAME = "_models_py3"
PAGED_NAME = "_paged_models"

def as_file_name(name):
    return name + ".py"

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):
    for models_module in models:
        models_path = models_module.__path__[0]
        _LOGGER.info(f"Working on {models_path}")
        if Path(models_path, as_file_name(MODEL_PY3_NAME)).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)

            # Swith the files
            shutil.rmtree(models_path)
            shutil.move(final_models_path, models_path)

def solve_one_model(models_module, output_folder):
    """Will build the compacted models in the output_folder"""

    models_classes = [
        (len(model_class.__mro__), model_name, 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__
    ]
    # Sort on MRO size first, and then alphabetically
    models_classes.sort(key=lambda x: (x[0], x[1]))

    # Just need the name of exceptions
    exceptions_classes = [
        model_name for model_name, model_class in vars(models_module).items()
        if model_name[0].isupper() and HttpOperationError in model_class.__mro__
    ]

    py2_models_classes = [
        (len_mro, model_name, path.replace("_py3.py", ".py"), None)
        for len_mro, model_name, path, _ in models_classes
    ]

    paged_models_classes = [
        (model_name, 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 = [
        (model_name, 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:
        # Can't be more than one enum file
        enum_file = Path(enum_models_classes[0][1])
        enum_file_module_name = "_"+enum_file.with_suffix('').name
        shutil.copyfile(enum_file, Path(output_folder, as_file_name(enum_file_module_name)))
    else:
        enum_file_module_name = None

    write_model_file(Path(output_folder, as_file_name(MODEL_PY3_NAME)), models_classes)
    write_model_file(Path(output_folder, as_file_name(MODEL_PY2_NAME)), py2_models_classes)
    write_paging_file(Path(output_folder, as_file_name(PAGED_NAME)), paged_models_classes)
    write_complete_init(
        Path(output_folder, "__init__.py"),
        models_classes,
        exceptions_classes,
        paged_models_classes,
        enum_models_classes,
        enum_file_module_name
    )

def write_model_file(output_file_path, classes_to_write):
    with open(output_file_path, "bw") as write_fd:
        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 write_complete_init(output_file_path, models, exceptions_classes, paging_models, enum_models, enum_file_module_name):
    with open(output_file_path, "bw") as write_fd:
        write_fd.write(copyright_header)

        write_fd.write(b"\ntry:\n")
        # Write py3 import
        for _, model_name, _, _ in models:
            write_fd.write(f"    from .{MODEL_PY3_NAME} import {model_name}\n".encode("utf-8"))
        # Py 3 exceptions
        for model_name in exceptions_classes:
            write_fd.write(f"    from .{MODEL_PY3_NAME} import {model_name}\n".encode("utf-8"))

        write_fd.write(b"except (SyntaxError, ImportError):\n")

        # Write py2 import
        for _, model_name, _, _ in models:
            write_fd.write(f"    from .{MODEL_PY2_NAME} import {model_name}\n".encode("utf-8"))
        # Py 2 exceptions
        for model_name in exceptions_classes:
            write_fd.write(f"    from .{MODEL_PY2_NAME} import {model_name}\n".encode("utf-8"))

        # Write paged import
        for model_name, _, _ in paging_models:
            write_fd.write(f"from .{PAGED_NAME} import {model_name}\n".encode("utf-8"))

        if enum_models:
            # Write enum model import
            for model_name, _, _ in enum_models:
                write_fd.write(f"from .{enum_file_module_name} import {model_name}\n".encode("utf-8"))

        write_fd.write(b"\n\n__all__=[\n")
        # Write all classes name
        for _, model_name, _, _ in models:
            write_fd.write(f"    '{model_name}',\n".encode("utf-8"))
        for model_name in exceptions_classes:
            write_fd.write(f"    '{model_name}',\n".encode("utf-8"))
        for model_name, _, _ in paging_models:
            write_fd.write(f"    '{model_name}',\n".encode("utf-8"))
        if enum_models:
            for model_name, _, _ in enum_models:
                write_fd.write(f"    '{model_name}',\n".encode("utf-8"))

        write_fd.write(b"]\n")


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
        ]


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(f"Looking for Autorest generated package in {module_prefix}")
    result = []
    try:
        _LOGGER.debug(f"Try {module_prefix}")
        importlib.import_module(".models", module_prefix)
        # If not exception, we found it
        _LOGGER.info(f"Found {module_prefix}")
        result.append(module_prefix)
    except ModuleNotFoundError:
        # No model, might dig deeper
        prefix_module = importlib.import_module(module_prefix)
        for _, sub_package, ispkg in pkgutil.iter_modules(prefix_module.__path__, module_prefix+"."):
            if ispkg:
                result += find_autorest_generated_folder(sub_package)
    return result


def main(prefix="azure.mgmt"):
    for autorest_package in find_autorest_generated_folder(prefix):
        models_module = find_models_to_change(autorest_package)
        solve_mro(models_module)

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    main(sys.argv[1] if len(sys.argv) >= 2 else "azure.mgmt")
	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
	"""

	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):
	for models_module in models:
	models_path = models_module.__path__[0]
	_LOGGER.info(f"Working on {models_path}")
	if 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)

	# Swith the files
	shutil.rmtree(models_path)
	shutil.move(final_models_path, models_path)

	def solve_one_model(models_module, output_folder):
	"""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])

	py2_models_classes = [
	(len_mro, path.replace("_py3.py", ".py"), None)
	for len_mro, path, _ in models_classes
	]

	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)
	write_model_file(Path(output_folder, "models.py"), py2_models_classes)
	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):
	with open(output_file_path, "bw") as write_fd:
	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
	]

	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(f"Looking for Autorest generated package in {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(f"Try {sub_package}")
	importlib.import_module(".models", sub_package)
	# If not exception, we found it
	_LOGGER.info(f"Found {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)

	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)
	solve_mro(models)