05-Custom_ASDF_Extensions/05-Custom-Extensions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asdf\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6 - Creating Custom ASDF Extensions\n",
    "\n",
    "Often we want to be able to save our \"custom\" python objects to ASDF in a \"seamless\"\n",
    "fashion. Earlier we were able to save various `astropy` objects using `asdf-astropy`\n",
    "in [tutorial 3](../03-Creating_ASDF_Files/03-Creating_ASDF_Files.ipynb). Here we will\n",
    "discuss how to create the necessary ASDF extension(s) to support doing this for\n",
    "a \"custom\" object so that ASDF can read (deserialize) and write (serialize) that\n",
    "object. Note that for a given object, we typically expect an ASDF extension\n",
    "supporting that object to \"round-trip\" that object, meaning the object can be serialized\n",
    "to ASDF and then deserialized from ASDF to an object which is an exact functional\n",
    "copy of the original object."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example Object\n",
    "\n",
    "Let's create a relatively simple Python object, which we would like to handle seamlessly with ASDF. For our purposes lets consider a geometric ellipse described by its\n",
    "- semi-major axis\n",
    "- semi-minor axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Ellipse:\n",
    "    \"\"\"An ellipse defined by semi-major and semi-minor axes.\n",
    "\n",
    "    Note: Using a dataclass to define the object so that we get `==` for free.\n",
    "    \"\"\"\n",
    "\n",
    "    semi_major: float\n",
    "    semi_minor: float"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that ASDF will handle objects contained inside the objects you wish to serialize provided that those objects are handled intrinsically by ASDF or an extension which handles each particular object is available for ASDF to use. For example, if we wanted to specify the axes of the ellipse using astropy Quantity objects (to attach units), so long as asdf-astropy is installed, ASDF would handle this transparently."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to Writing an Extension\n",
    "\n",
    "An ASDF extension requires two components to function properly:\n",
    "\n",
    "1. A `tag` for the object, so that ASDF identify/validate the object to deserialize it from an ASDF\n",
    "file.\n",
    "2. A `Converter` for the object, so ASDF knows how to serialize and deserialize the object to and from\n",
    "an ASDF file.\n",
    "\n",
    "The `tag` is defined through the schemas and related resources for ASDF to use. While the `Converter`\n",
    "is a python object which provides the code the ASDF library executes in order to handle the serialization\n",
    "or deserialization process. The `Converter` is then wrapped inside an ASDF `Extension` object (which\n",
    "can contain several different `Converter`s), which is then added to ASDF (typically) via an entry\n",
    "point."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating a `tag`\n",
    "\n",
    "Recall that ASDF supports the use of schemas for validating the correctness of the information stored within its files.\n",
    "Often one wishes to create a schema for a specific object so that the particular object the schema description can be\n",
    "reused in other schemas. A `tag` is a reference to a specific schema or set of schemas that a particular value in an ASDF\n",
    "file tree need to satisfy. This `tag` is then used within the `yaml` metadata to identify the sub-tree which represents\n",
    "the object within the ASDF file. Thus the tag serves two purposes:\n",
    "\n",
    "1. Identifying the schema used to validate a sub-tree of the ASDF tree.\n",
    "2. Identifying the object a particular sub-tree describes.\n",
    "\n",
    "This means that in order to create a `tag` for a given Python object we really need to create resource `yaml` files for ASDF\n",
    "to do two things:\n",
    "\n",
    "1. Contain schema(s) used by that `tag`.\n",
    "2. Create an association between the schema(s) and the `tag`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating a Schema\n",
    "\n",
    "Recall that in tutorial 4, we discussed in depth how to create ASDF schemas. In particular note that schemas are typically\n",
    "stored in `yaml` files which are then loaded into ASDF via an entry point.\n",
    "\n",
    "To begin with lets create the schema for `Ellipse` dynamically (without needing to package our code for an entry point):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ellipse_uri = \"asdf://example.com/example-project/schemas/ellipse-1.0.0\"\n",
    "\n",
    "ellipse_schema_content = f\"\"\"\n",
    "%YAML 1.1\n",
    "---\n",
    "$schema: http://stsci.edu/schemas/yaml-schema/draft-01\n",
    "id: {ellipse_uri}\n",
    "\n",
    "type: object\n",
    "properties:\n",
    "  semi_major:\n",
    "    type: number\n",
    "  semi_minor:\n",
    "    type: number\n",
    "required: [semi_major, semi_minor]\n",
    "...\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This can then be dynamically added to ASDF using the `add_resource_mapping`. This adds a map (`dict`) between a `uri`\n",
    "(universal resource identifier) string and the content of the resource to ASDF. Later, when working with the entry\n",
    "points directly, we will need to specify how to build these mappings.\n",
    "\n",
    "Note we highly recommend as best practice to always have the `id `for any resource and the `uri` string be the same. This\n",
    "is to limit the possibility of confusing how to look-up the given schema as JSON schema (the base language/library used\n",
    "for ASDF schemas) uses the `id` field to reference resources among one-another, while ASDF uses the `uri` as keys to find\n",
    "those resources on disk. One does not have to follow this practice, but it is highly discouraged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asdf.get_config().add_resource_mapping({ellipse_uri: ellipse_schema_content})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets now load and check that the schema we just created is a valid schema:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = asdf.schema.load_schema(ellipse_uri)\n",
    "asdf.schema.check_schema(schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `asdf.schema.check_schema` will work directly on any `yaml` file loaded through the `pyyaml` interface.\n",
    "Lets also attempt to validate a portion of an ASDF tree for `Ellipse` against this schema:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Valid tree\n",
    "\n",
    "test_ellipse_object = {\"semi_major\": 1.0, \"semi_minor\": 2.0}\n",
    "asdf.schema.validate(test_ellipse_object, schema=schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "# Invalid tree\n",
    "\n",
    "test_ellipse_object = {\"semi_major\": 3.0}\n",
    "asdf.schema.validate(test_ellipse_object, schema=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that ASDF provides a `pytest` plugin which can be configured to automatically generate unit tests which will\n",
    "check and validate all of the schemas in your package. In fact, you can include \"examples\" of the tree the schema\n",
    "is checking and the plugin will test that those examples do correctly validate against the schema itself."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Exercise 1\n",
    "\n",
    "Create, add, and check a schema for the `Rectangle` object below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Rectangle:\n",
    "    base: float\n",
    "    height: float"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating the `tag` Itself\n",
    "\n",
    "ASDF uses a special `schema` to specify the `tag`s for a given ASDF extension. This special type of `schema` is called a\n",
    "`manifest` which lists each `tag` as a pair of `uris`:\n",
    "\n",
    "- `tag_uri`, the `uri` which will be used for the `tag`.\n",
    "- `schema_uri`, the ASDF `uri` used to reference the specific `schema`(s) involved.\n",
    "\n",
    "This allows for a given `schema` to be reused for multiple `tag`s. Such as for objects which contain the same serializable\n",
    "data, but have different Python functionalities which need to be distinguished.\n",
    "\n",
    "The following is an example for creating/adding a manifest for an extension which has the resources for the `Ellipse` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes_manifest_uri = \"asdf://example.com/example-project/manifests/shapes-1.0.0\"\n",
    "shapes_extension_uri = \"asdf://example.com/example-project/extensions/shapes-1.0.0\"\n",
    "ellipse_tag = \"asdf://example.com/example-project/tags/ellipse-1.0.0\"\n",
    "\n",
    "shapes_manifest_content = f\"\"\"\n",
    "%YAML 1.1\n",
    "---\n",
    "id: {shapes_manifest_uri}\n",
    "extension_uri: {shapes_extension_uri}\n",
    "\n",
    "title: Example Shape extension 1.0.0\n",
    "description: Tags for example shape objects.\n",
    "\n",
    "tags:\n",
    "  - tag_uri: {ellipse_tag}\n",
    "    schema_uri: {ellipse_uri}\n",
    "...\n",
    "\"\"\"\n",
    "\n",
    "asdf.get_config().add_resource_mapping({shapes_manifest_uri: shapes_manifest_content})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We again add this `manifest` to ASDF via the `add_resource_mapping` interface. Note that the `extension_uri` field defines\n",
    "the `uri` that the whole `Extension` (resource(s) combined with Converter(s)) uses within ASDF. The `extension_uri` will\n",
    "be referenced later by the `Extension` object so that the extension code will be available when the correct resources are\n",
    "available and vice-versa.\n",
    "\n",
    "Since the `manifest` is just like any other schema we can check it in the same way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check\n",
    "schema = asdf.schema.load_schema(shapes_manifest_uri)\n",
    "asdf.schema.check_schema(schema)\n",
    "asdf.schema.validate(shapes_manifest_content, schema=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Exercise 2\n",
    "\n",
    "Add a `tag` for your `rectangle` schema to the `shapes-1.0.0` manifest and add your manifest to ASDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a `Converter`\n",
    "\n",
    "All converters should be constructed as subclasses of the abstract type `asdf.extension.Converter`,\n",
    "which requires that you define two methods:\n",
    "\n",
    "1. `to_yaml_tree`: which converts a Python object into an ASDF tree.\n",
    "2. `from_yaml_tree`: which converts an ASDF tree into a python object.\n",
    "\n",
    "Note that these methods can account for the type/tag of the objects attempting to be converted.\n",
    "\n",
    "Moreover your converter also needs to define the following two variables:\n",
    "\n",
    "1. `tags`: A list of tags that this converter will use when reading ASDF.\n",
    "2. `types`: A list of Python (object) types that this converter will use when writing ASDF.\n",
    "\n",
    "Note that these lists do not need to be indexed with respect to each other, and that in order for\n",
    "the converter to actually be used by ASDF, at least one of the `tags` needs to be registered as a\n",
    "resource with ASDF (usually via the entry point).\n",
    "\n",
    "An example converter for `Ellipse`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EllipseConverter(asdf.extension.Converter):\n",
    "    tags = [ellipse_tag]\n",
    "    types = [Ellipse]\n",
    "\n",
    "    def to_yaml_tree(self, obj, tag, ctx):\n",
    "        return {\n",
    "            \"semi_major\": obj.semi_major,\n",
    "            \"semi_minor\": obj.semi_minor,\n",
    "        }\n",
    "\n",
    "    def from_yaml_tree(self, node, tag, ctx):\n",
    "        return Ellipse(semi_major=node[\"semi_major\"], semi_minor=node[\"semi_minor\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that the converter itself will be added to ASDF via an `Extenstion` not as a resource like the\n",
    "`schema`s above.  Note that, for performance of the entry points, one will normally defer the `import`\n",
    "of the object (`Ellipse` in this case) until `from_yaml_tree` is actually called."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 3\n",
    "\n",
    "Create a converter for the `Rectangle` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the Full Extension\n",
    "\n",
    "Now lets dynamically create an extension for ASDF to support the `Ellipse` object using the `EllipseConverter` we just\n",
    "created and the `ellipse_tag` we created earlier.\n",
    "\n",
    "This can be accomplished via using the `asdf.extensions.ManifestExtension.from_uri` constructor, which in our case requires\n",
    "two arguments:\n",
    "\n",
    "1. The `manifest_uri`, the `uri` the manifest was added under.\n",
    "2. The `converters`, a list of instances of `Converter` classes.\n",
    "\n",
    "Note that one can also pass a list of `Compressor` objects (ASDF objects to handle custom binary block compression).\n",
    "\n",
    "An instance of the extension object can then be dynamically added to ASDF using the `add_extension` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes_extension = asdf.extension.ManifestExtension.from_uri(\n",
    "    shapes_manifest_uri, converters=[EllipseConverter()]\n",
    ")\n",
    "asdf.get_config().add_extension(shapes_extension)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternately one can create an `Extension` object directly by extending `asdf.extension.Extension` and\n",
    "specifying the variables:\n",
    "\n",
    "1. `extension_uri`, the `extension_uri` specified within the `manifest` in question.\n",
    "2. `converters`, the list of `Converter` objects making up the extension.\n",
    "3. `tags`, the list of tags those `Converter` objects use from the `manifest` referenced.\n",
    "\n",
    "which can then be dynamically added in exactly the same way.\n",
    "\n",
    "Note that the `from_uri` constructor, figures out all this information from the `uri` and\n",
    "`Converter` objects themselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EllipseExtension(asdf.extension.Extension):\n",
    "    extension_uri = shapes_extension_uri\n",
    "    converters = [EllipseConverter()]\n",
    "    tags = [ellipse_tag]\n",
    "\n",
    "\n",
    "asdf.get_config().add_extension(EllipseExtension())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Testing the `Ellipse` Extension\n",
    "\n",
    "Lets now check that we can round-trip an `Ellipse` object through ASDF:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ellipse = Ellipse(1.0, 2.0)\n",
    "\n",
    "with asdf.AsdfFile() as af:\n",
    "    af[\"ellipse\"] = ellipse\n",
    "    af.write_to(\"ellipse.asdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's examine the contents of the ASDF file and then read/compare them to our original object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"ellipse.asdf\") as f:\n",
    "    print(f.read())\n",
    "\n",
    "with asdf.open(\"ellipse.asdf\") as af:\n",
    "    print(af[\"ellipse\"])\n",
    "    assert af[\"ellipse\"] == ellipse"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Exercise 4\n",
    "\n",
    "Create, add, and test an extension for your `RectangleConverter`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Entry-Points to Automatically Extend ASDF\n",
    "\n",
    "Obviously, having to dynamically add all the resources and extensions to ASDF every time you want to work with a custom\n",
    "object is tedious. Indeed, `asdf-astropy` only needs to be installed so that its `Extensions` are available for ASDF to\n",
    "use. This is accomplished by using Python entry-points (mechanism for one python package to communicate information to\n",
    "another Python package), to enable automatic discovery and loading of resources and extensions for ASDF.\n",
    "\n",
    "Since entry-points are a means for Python packages to communicate with one-another, their use requires you to package\n",
    "your Python code, which is can be a complex issue. Thus we will assume that you have an existing Python package, that\n",
    "you wish to add our example ASDF extension to.\n",
    "\n",
    "To create our entry-points we will need to make three modifications to the packaging components of the existing Python\n",
    "package:\n",
    "\n",
    "1. Create an entry point to add the resources to ASDF.\n",
    "2. Create an entry point to add the extensions to ASDF.\n",
    "\n",
    "Note that we will assume that you are using the `setup.cfg` file to configure your python package. This can also be\n",
    "done via the `pyproject.toml` file in a similar fashion (see ASDF docs for details)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an Entry-Point for the Resources\n",
    "\n",
    "ASDF treats the information it receives from the entry-points it checks for resources as a function that it can evaluate\n",
    "to get a list of resource mappings. To begin suppose that there your package is called `asdf_shapes` and the function you\n",
    "need to call in order to get this list of mappings is `called get_resource_mappings` and is located in the `integration`\n",
    "module, that is you need to import `get_resource_mappings` from `asdf_shapes.integration`. Given this setup you will need\n",
    "to add the following to your `setup.cfg`:\n",
    "\n",
    "```\n",
    "[options.entry_points]\n",
    "asdf.resource_mappings =\n",
    "    asdf_shapes_schemas = asdf_shapes.integration:get_resource_mappings\n",
    "```\n",
    "\n",
    "Breaking this down:\n",
    "\n",
    "- The entry-point ASDF checks for resources is `asdf.resource_mappings`.\n",
    "- The identifier for your package's resources in this case is `asdf_shapes_schemas`.\n",
    "- The method to that needs to be executed is `asdf_shapes.integration:get_resource_mappings` which corresponds to the form\n",
    "`module:function`.\n",
    "\n",
    "Now lets talk about how to create the `get_resource_mappings` function. First, lets go ahead a create the `yaml` files\n",
    "for the resources we used in our example in order to illustrate an example organization of these resource files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_root = \"resources/schemas\"\n",
    "manifest_root = \"resources/manifests\"\n",
    "\n",
    "os.makedirs(schema_root, exist_ok=True)\n",
    "os.makedirs(manifest_root, exist_ok=True)\n",
    "\n",
    "with open(f\"{schema_root}/ellipse-1.0.0.yaml\", \"w\") as f:\n",
    "    f.write(ellipse_schema_content)\n",
    "\n",
    "with open(f\"{manifest_root}/shapes-1.0.0.yaml\", \"w\") as f:\n",
    "    f.write(shapes_manifest_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Normally we organize the resource files into a directory structure (as we just did) which can be parsed to form part of\n",
    "the URI (`id`) used for each resource document as part of the file path. This is done so that adding resources can be\n",
    "performed by ASDF by crawling these directory structures where the directory structure helps to determine the `uri`.\n",
    "\n",
    "ASDF provides the `asdf.resource.DirectoryResourceMapping` object to crawl resource directories.  It allows us to turn\n",
    "these directory structures into resource mappings, which can subsequently be added to ASDF using the entry-points.\n",
    "\n",
    "These objects require two input parameters:\n",
    "\n",
    "1. A path to the root directory which contains the resources to be added.\n",
    "2. The prefix that will be used together with the file names to generate the `uri` for the resource in question.\n",
    "\n",
    "There are some optional inputs:\n",
    "1. `recursive`: (default `False`) which determines if the object will search recursively through subdirectories.\n",
    "2. `filename_pattern`: (default: `*.yaml`) Glob pattern for the files that should be added.\n",
    "3. `stem_filename`: (default: `True`) determine if the file extension should be removed when creating the `uri`.\n",
    "\n",
    "In this case we do not need to set any of the file, here we need to do only the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In module asdf_shapes.integration\n",
    "def get_resource_mappings():\n",
    "    schema_prefix = \"asdf://example.com/example-project/schemas/\"\n",
    "    manifest_prefix = \"asdf://example.com/example-project/manifests/\"\n",
    "    return [\n",
    "        asdf.resource.DirectoryResourceMapping(schema_root, schema_prefix),\n",
    "        asdf.resource.DirectoryResourceMapping(manifest_root, manifest_prefix),\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Which can then be referenced by the entry-point. Note that for performance reasons, we suggest you limit the top-level\n",
    "imports of the file(s) you load your entry points from to as few as possible, going as far as deferring imports to inside\n",
    "the entry-point functions when possible. This is because asdf will import all of these models immediately when `asdf.open`\n",
    "is called meaning large imports will cause noticeable delays especially when using the command-line interface."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an Entry-Point for the Extensions\n",
    "\n",
    "In a similar fashion to resources, ASDF assumes the entry-points it checks for extensions as functions which return\n",
    "lists of `asdf.extension.Extension` objects. Thus lets assume your function is called `get_extensions` and is in the\n",
    "`asdf_shapes.integration` module alongside get_resource_mappings. Adding the entry-point to `setup.cfg` for this would\n",
    "look something like:\n",
    "\n",
    "```\n",
    "[options.entry_points]\n",
    "asdf.extensions =\n",
    "    asdf_shapes_extensions = asdf_shapes.integration:get_extensions\n",
    "```\n",
    "\n",
    "Breaking this down:\n",
    "\n",
    "- The entry-point that ASDF checks for `Extension`s is `asdf.extensions`.\n",
    "- The identifier for your package's extensions in this case is `asdf_shapes_extensions`.\n",
    "- The method to that needs to be executed is `asdf_shapes.integration:get_extensions` which corresponds to the form\n",
    "`module:function`.\n",
    "\n",
    "The structure of `get_extensions` will be similar to that for `get_resource_mappings`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In module asdf_shapes.integration\n",
    "def get_extensions():\n",
    "    # import EllipseConverter inside this function\n",
    "    return [\n",
    "        asdf.extension.ManifestExtension.from_uri(\n",
    "            shapes_manifest_uri, converters=[EllipseConverter()]\n",
    "        )\n",
    "    ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once your package is installed with these changes, ASDF will automatically detect and use your ASDF extension as needed\n",
    "in a seamless fashion."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extending Our Example Object\n",
    "\n",
    "Suppose that we want to extend our object so that it represents an Ellipse rotated around the origin off the XY plane,\n",
    "that is add a position angle:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class RotatedEllipse(Ellipse):\n",
    "    position_angle: float"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extend the Schema\n",
    "\n",
    "JSON schema does not support the concept of inheritance, which makes \"extending\" an existing schema somewhat awkward.\n",
    "What we do instead is create a schema which adds attributes to the existing schema via the `allOf` operation. In this case,\n",
    "we can define the a schema for `RotatedEllipse` by adding a `position_angle` property:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rotated_ellipse_uri = \"asdf://example.com/example-project/schemas/rotated_ellipse-1.0.0\"\n",
    "\n",
    "rotated_ellipse_schema_content = f\"\"\"\n",
    "%YAML 1.1\n",
    "---\n",
    "$schema: http://stsci.edu/schemas/yaml-schema/draft-01\n",
    "id: {rotated_ellipse_uri}\n",
    "\n",
    "allOf:\n",
    "  - $ref: {ellipse_uri}\n",
    "  - properties:\n",
    "      position_angle:\n",
    "        type: number\n",
    "    required: [position_angle]\n",
    "...\n",
    "\"\"\"\n",
    "\n",
    "asdf.get_config().add_resource_mapping(\n",
    "    {rotated_ellipse_uri: rotated_ellipse_schema_content}\n",
    ")\n",
    "\n",
    "# check\n",
    "schema = asdf.schema.load_schema(rotated_ellipse_uri)\n",
    "asdf.schema.check_schema(schema)\n",
    "\n",
    "test_rotated_ellipse_object = {\n",
    "    \"semi_major\": 1.0,\n",
    "    \"semi_minor\": 2.0,\n",
    "    \"position_angle\": 3.0,\n",
    "}\n",
    "asdf.schema.validate(test_rotated_ellipse_object, schema=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 5\n",
    "\n",
    "Create a schema for the `RectanglularPrism` object below using the one for `Rectangle`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class RectangularPrism(Rectangle):\n",
    "    depth: float"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an Updated Manifest\n",
    "\n",
    "Now lets extend the `shapes-1.0.0` manifest to include a `rotated_ellipse-1.0.0`\n",
    "tag. Note that if a manifest is already released and in use, it is recommended that\n",
    "one create a new manifest whenever schemas or tags need to be modified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rotated_ellipse_tag = \"asdf://example.com/example-project/tags/rotated_ellipse-1.0.0\"\n",
    "\n",
    "shapes_manifest_content = f\"\"\"\n",
    "%YAML 1.1\n",
    "---\n",
    "id: {shapes_manifest_uri}\n",
    "extension_uri: {shapes_extension_uri}\n",
    "\n",
    "title: Example Shape extension 1.0.0\n",
    "description: Tags for example shape objects.\n",
    "\n",
    "tags:\n",
    "  - tag_uri: {ellipse_tag}\n",
    "    schema_uri: {ellipse_uri}\n",
    "\n",
    "  - tag_uri: {rotated_ellipse_tag}\n",
    "    schema_uri: {rotated_ellipse_uri}\n",
    "...\n",
    "\"\"\"\n",
    "\n",
    "asdf.get_config().add_resource_mapping({shapes_manifest_uri: shapes_manifest_content})\n",
    "\n",
    "\n",
    "# check\n",
    "schema = asdf.schema.load_schema(shapes_manifest_uri)\n",
    "asdf.schema.check_schema(schema)\n",
    "asdf.schema.validate(shapes_manifest_content, schema=schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Exercise 6\n",
    "\n",
    "Update the manifest with a tag for `RectangularPrism`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an Updated `Converter`\n",
    "\n",
    "The \"simplest\" approach to creating a `Converter` for `RotatedEllipse` would be to simply\n",
    "create a new converter as we did above for `Ellipse`; however, we can also take advantage\n",
    "of the fact that multiple `tags` and `types` can be listed. Note that when multiple tags\n",
    "are handled by the same `Converter`, we need to also implement a `select_tag` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RotatedEllipseConverter(asdf.extension.Converter):\n",
    "    tags = [ellipse_tag, rotated_ellipse_tag]\n",
    "    types = [Ellipse, RotatedEllipse]\n",
    "\n",
    "    def select_tag(self, obj, tag, ctx):\n",
    "        if isinstance(obj, RotatedEllipse):\n",
    "            return rotated_ellipse_tag\n",
    "        elif isinstance(obj, Ellipse):\n",
    "            return ellipse_tag\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown object {type(obj)}\")\n",
    "\n",
    "    def to_yaml_tree(self, obj, tag, ctx):\n",
    "        tree = {\n",
    "            \"semi_major\": obj.semi_major,\n",
    "            \"semi_minor\": obj.semi_minor,\n",
    "        }\n",
    "\n",
    "        if tag == rotated_ellipse_tag:\n",
    "            tree[\"position_angle\"] = obj.position_angle\n",
    "\n",
    "        return tree\n",
    "\n",
    "    def from_yaml_tree(self, node, tag, ctx):\n",
    "        if tag == ellipse_tag:\n",
    "            return Ellipse(**node)\n",
    "        elif tag == rotated_ellipse_tag:\n",
    "            return RotatedEllipse(**node)\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown tag {tag}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 7\n",
    "\n",
    "Create a converter to support `Rectangle` and `RectagularPrism`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating an Updated `Extension`\n",
    "\n",
    "We can now use this converter to create a new \"updated\" extension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes_extension = asdf.extension.ManifestExtension.from_uri(\n",
    "    shapes_manifest_uri, converters=[RotatedEllipseConverter()]\n",
    ")\n",
    "asdf.get_config().add_extension(shapes_extension)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checking the New Extension\n",
    "\n",
    "Lets check this new extension by writing both an `Ellipse` and `Ellipse3D` object\n",
    "to ASDF:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ellipse = Ellipse(1.0, 2.0)\n",
    "rotated_ellipse = RotatedEllipse(1.0, 2.0, 3.0)\n",
    "\n",
    "with asdf.AsdfFile() as af:\n",
    "    af[\"ellipse\"] = ellipse\n",
    "    af[\"rotated_ellipse\"] = rotated_ellipse\n",
    "    af.write_to(\"rotated_ellipse.asdf\")\n",
    "\n",
    "# Check\n",
    "with open(\"rotated_ellipse.asdf\") as f:\n",
    "    print(f.read())\n",
    "\n",
    "with asdf.open(\"rotated_ellipse.asdf\") as af:\n",
    "    print(af[\"ellipse\"])\n",
    "    assert af[\"ellipse\"] == ellipse\n",
    "\n",
    "    print(af[\"rotated_ellipse\"])\n",
    "    assert af[\"rotated_ellipse\"] == rotated_ellipse"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Exercise 8\n",
    "\n",
    "Create and check an extension supporting `Ellipse`, `RotatedEllipse`, `Rectangle`, and `RectangularPrism`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "c54ecae8f3b9a546603254558bcb3ca9dea5c7af451a66a7914c053836710349"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}