typechecks

llamazure/azrest/models.py

@@ -9,6 +9,7 @@
 from pydantic import BaseModel, BeforeValidator, Field
 
 Ret_T = TypeVar("Ret_T")
+Ret_T0 = TypeVar("Ret_T0")
 ReadOnly = Optional[Ret_T]
 
 
@@ -61,9 +62,9 @@ def add_params(self, params: Dict[str, str]) -> Req:
 	def add_param(self, name: str, value: str) -> Req:
 		return dataclasses.replace(self, params={**self.params, **{name: value}})
 
-	def with_ret_t(self, ret_t: Type[Ret_T]) -> Req:
+	def with_ret_t(self, ret_t: Type[Ret_T0]) -> Req[Ret_T0]:
 		"""Override the return type"""
-		return dataclasses.replace(self, ret_t=ret_t)
+		return dataclasses.replace(self, ret_t=ret_t)  # type: ignore
 
 
 @dataclass
@@ -156,6 +157,20 @@ class AzureErrorAdditionInfo(BaseModel):
 
 
 def ensure(a: Optional[T]) -> T:
+	"""Ensure the result is not None"""
 	if a is None:
 		raise TypeError("value was None")
 	return a
+
+
+P0 = TypeVar("P0", bound=BaseModel)
+P1 = TypeVar("P1", bound=BaseModel)
+
+
+def cast_as(obj: P0, cls: Type[P1]) -> P1:
+	"""
+	Cast one model into another.
+
+	Useful for turning a Foo into a FooUpdateParameters.
+	"""
+	return cls.model_validate(obj.model_dump())

llamazure_tools/migrate/dashboard.py

@@ -1,5 +1,4 @@
 """Migrate an Azure Dashboard to a different Log Analytics Workspace"""
-import dataclasses
 import json
 from dataclasses import dataclass
 from datetime import datetime
@@ -9,10 +8,11 @@
 from azure.identity import DefaultAzureCredential
 
 from llamazure.azrest.azrest import AzRest
+from llamazure.azrest.models import cast_as
 from llamazure.rid import rid
 from llamazure.rid.rid import Resource
 from llamazure_tools.migrate.portal.r.m.portal.portal import AzDashboards, Dashboard, PatchableDashboard  # pylint: disable=E0611,E0401
-from llamazure_tools.migrate.util import JSONTraverser
+from llamazure_tools.migrate.util import JSONTraverser, rid_params
 
 
 @dataclass
@@ -35,7 +35,7 @@ def migrate(self):
 
 	def get_dashboard(self) -> dict:
 		"""Retrieve the current dashboard data from Azure."""
-		return self.az.call(dataclasses.replace(AzDashboards.Get(self.dashboard.sub.uuid, self.dashboard.rg.name, self.dashboard.name), ret_t=dict))
+		return self.az.call(AzDashboards.Get(*rid_params(self.dashboard)).with_ret_t(dict))
 
 	def transform(self, dashboard: dict) -> dict:
 		"""Transform the dashboard data using the provided transformer."""
@@ -44,12 +44,9 @@ def transform(self, dashboard: dict) -> dict:
 	def put_dashboard(self, transformed: dict):
 		"""Update the dashboard in Azure with the transformed data."""
 		d = Dashboard(**transformed)
-		p = PatchableDashboard(
-			properties=d.properties.model_dump(),
-			tags=d.tags,
-		)
+		p = cast_as(d, PatchableDashboard)
 		self.az.call(
-			AzDashboards.Update(self.dashboard.sub.uuid, self.dashboard.rg.name, self.dashboard.name, p),
+			AzDashboards.Update(*rid_params(self.dashboard), p),
 		)
 
 	def make_backup(self, dashboard: dict):
@@ -68,7 +65,9 @@ def migrate(resource_id: str, replacements: str, backup_directory: str):
 	az = AzRest.from_credential(DefaultAzureCredential())
 
 	replacements = json.loads(replacements)
+	assert isinstance(replacements, dict)
 	resource = rid.parse(resource_id)
+	assert isinstance(resource, rid.Resource)
 	transformer = JSONTraverser(replacements)
 	migrator = Migrator(az, resource, transformer, Path(backup_directory))
 

llamazure_tools/migrate/util.py

@@ -1,6 +1,9 @@
 """Utils"""
 from dataclasses import dataclass
-from typing import Any, Dict
+from typing import Any, Dict, Tuple
+
+from llamazure.azrest.models import ensure
+from llamazure.rid import rid
 
 
 @dataclass
@@ -19,3 +22,7 @@ def traverse(self, obj: Any) -> Any:
 			return self.replacements.get(obj, obj)
 		else:
 			return obj
+
+
+def rid_params(res: rid.Resource) -> Tuple[str, str, str]:
+	return ensure(res.sub).uuid, ensure(res.rg).name, res.name

llamazure_tools/migrate/workbook.py

@@ -8,10 +8,11 @@
 from azure.identity import DefaultAzureCredential
 
 from llamazure.azrest.azrest import AzRest
+from llamazure.azrest.models import cast_as
 from llamazure.rid import rid
 from llamazure.rid.rid import Resource
-from llamazure_tools.migrate.applicationinsights.r.m.insights.workbooks import AzWorkbooks, Workbook  # pylint: disable=E0611,E0401
-from llamazure_tools.migrate.util import JSONTraverser
+from llamazure_tools.migrate.applicationinsights.r.m.insights.workbooks import AzWorkbooks, Workbook, WorkbookUpdateParameters  # pylint: disable=E0611,E0401
+from llamazure_tools.migrate.util import JSONTraverser, rid_params
 
 
 @dataclass
@@ -33,7 +34,7 @@ def migrate(self):
 
 	def get_workbook(self) -> Workbook:
 		"""Retrieve the current workbook data from Azure."""
-		return self.az.call(AzWorkbooks.Get(self.workbook.sub.uuid, self.workbook.rg.name, self.workbook.name, canFetchContent=True))
+		return self.az.call(AzWorkbooks.Get(*rid_params(self.workbook), canFetchContent=True))
 
 	def transform(self, workbook: Workbook) -> Workbook:
 		"""Transform the workbook data using the provided transformer."""
@@ -43,7 +44,7 @@ def transform(self, workbook: Workbook) -> Workbook:
 	def put_workbook(self, transformed: Workbook):
 		"""Update the dashboard in Azure with the transformed data."""
 		self.az.call(
-			AzWorkbooks.Update(self.workbook.sub.uuid, self.workbook.rg.name, self.workbook.name, transformed),
+			AzWorkbooks.Update(*rid_params(self.workbook), cast_as(transformed, WorkbookUpdateParameters)),
 		)
 
 	def make_backup(self, workbook: Workbook):
@@ -62,7 +63,9 @@ def migrate(resource_id: str, replacements: str, backup_directory: str):
 	az = AzRest.from_credential(DefaultAzureCredential())
 
 	replacements = json.loads(replacements)
+	assert isinstance(replacements, dict)
 	resource = rid.parse(resource_id)
+	assert isinstance(resource, rid.Resource)
 	transformer = JSONTraverser(replacements)
 	migrator = Migrator(az, resource, transformer, Path(backup_directory))