try make mypy happy

lambeq/bobcat/tree.py

@@ -292,7 +292,8 @@ def deps(self) -> list[Dependency]:
         return self.deps_and_tags[0]
 
     def to_json(self) -> Dict[str, Any]:
-        data = {'type': str(self.cat), 'rule': self.rule.name}
+        data: Dict[str, Any] = {
+            'type': str(self.cat), 'rule': self.rule.name}
         if self.left:
             data['children'] = [self.left.to_json()]
         if self.right:

lambeq/text2diagram/ccg_tree.py

@@ -25,6 +25,8 @@
 from lambeq.backend.grammar import Diagram, Id, Word
 from lambeq.text2diagram.ccg_rule import CCGRule
 from lambeq.text2diagram.ccg_type import CCGType
+from lambeq.text2diagram.depccg_parser import Tree as DepCCGTree
+from lambeq.bobcat import ParseTree as BobcatTree
 
 # Types
 _JSONDictT = Dict[str, Any]
@@ -44,6 +46,7 @@ def __init__(self,
                  rule: CCGRule | str = CCGRule.UNKNOWN,
                  biclosed_type: CCGType,
                  children: Iterable[CCGTree] | None = None,
+                 original: DepCCGTree | BobcatTree | None = None,
                  metadata: dict[Any, Any] | None = None) -> None:
         """Initialise a CCG tree.
 
@@ -68,6 +71,7 @@ def __init__(self,
         self.rule = CCGRule(rule)
         self.biclosed_type = biclosed_type
         self.children = list(children) if children is not None else []
+        self.original = original
         self.metadata = metadata if metadata is not None else {}
 
         n_children = len(self.children)
@@ -186,8 +190,8 @@ def without_trivial_unary_rules(self) -> CCGTree:
 
     def to_json(self, original=False) -> _JSONDictT:
         """Convert tree into JSON form."""
-        if original:
-            return self.metadata['original'].to_json()
+        if original and self.original is not None:
+            return self.original.to_json()
         if self is None:  # Allows doing CCGTree.to_json(X) for optional X
             return None  # type: ignore[unreachable]
 

lambeq/text2diagram/depccg_parser.py

@@ -25,7 +25,7 @@
 from collections.abc import Iterable
 import functools
 import logging
-from typing import Any, TYPE_CHECKING
+from typing import Dict, Any, TYPE_CHECKING
 
 from lambeq.core.globals import VerbosityLevel
 from lambeq.core.utils import (
@@ -42,6 +42,7 @@
     from depccg.annotator import (annotate_XX, english_annotator,
                                   japanese_annotator)
     from depccg.cat import Category
+    from depccg.tree import Tree
     from lambeq.backend.grammar import Diagram
 
 
@@ -468,8 +469,8 @@ def _build_ccgtree(tree: depccg.tree.Tree) -> CCGTree:
                        metadata={'original': tree})
 
 
-def depccg_to_json(tree: depccg.tree.Tree) -> dict[str, Any]:
-    data = {'type': str(tree.cat)}
+def depccg_to_json(tree: depccg.tree.Tree) -> Dict[str, Any]:
+    data: Dict[str, Any] = {'type': str(tree.cat)}
     if tree.is_leaf:
         data['rule'], data['text'] = 'L', tree.word
     else: