Fix issue processing dex treemap

src/launchpad/size/treemap/dex_element_builder.py

@@ -62,37 +62,47 @@ def _build_package_tree(self) -> list[TreemapElement]:
             # Add the class to the leaf package
             leaf_package = package_tree
             for package_part in package_parts:
-                leaf_package = leaf_package[package_part]["packages"]
-            leaf_package[class_name] = {"class_def": class_def}
+                if package_part not in leaf_package:
+                    leaf_package[package_part] = {"packages": {}, "classes": {}}
+                if package_part == package_parts[-1]:
+                    # This is the final package, add the class here
+                    leaf_package[package_part]["classes"][class_name] = {"class_def": class_def}
+                else:
+                    # Navigate to the next level
+                    leaf_package = leaf_package[package_part]["packages"]
 
         return self._convert_tree_to_elements(package_tree)
 
     def _convert_tree_to_elements(self, package_tree: dict[str, dict], parent_path: str = "") -> list[TreemapElement]:
         elements: list[TreemapElement] = []
 
         for name, node in package_tree.items():
-            if "class_def" in node:
-                class_def = node["class_def"]
-                class_element = self._create_class_element(class_def)
-                elements.append(class_element)
-            else:
-                package_path = f"{parent_path}.{name}" if parent_path else f"{name}"
-
-                # Process children (sub-packages and classes)
-                children = self._convert_tree_to_elements(node["packages"], package_path)
-
-                total_size = sum(child.size for child in children)
-
-                elements.append(
-                    TreemapElement(
-                        name=name,
-                        size=total_size,
-                        type=TreemapType.DEX,
-                        path=package_path,
-                        is_dir=True,
-                        children=children,
-                    )
+            package_path = f"{parent_path}.{name}" if parent_path else f"{name}"
+
+            # Process sub-packages
+            children = []
+            if "packages" in node:
+                children.extend(self._convert_tree_to_elements(node["packages"], package_path))
+
+            # Process classes in this package
+            if "classes" in node:
+                for class_name, class_node in node["classes"].items():
+                    class_def = class_node["class_def"]
+                    class_element = self._create_class_element(class_def)
+                    children.append(class_element)
+
+            total_size = sum(child.size for child in children)
+
+            elements.append(
+                TreemapElement(
+                    name=name,
+                    size=total_size,
+                    type=TreemapType.DEX,
+                    path=package_path,
+                    is_dir=True,
+                    children=children,
                 )
+            )
 
         return elements