diff --git a/examples/chess.d2 b/examples/chess.d2
index a1f4e36..8c2cedd 100644
--- a/examples/chess.d2
+++ b/examples/chess.d2
@@ -1,12 +1,7 @@
 # Actors
 hans: Hans Niemann
 
-defendants: |md
-  # Defendants
-
-  hello world
-| {
-  shape: text
+defendants: {
   mc: Magnus Carlsen
   playmagnus: Play Magnus Group
   chesscom: Chess.com
diff --git a/grammar.js b/grammar.js
index dcd8209..bbcc51e 100644
--- a/grammar.js
+++ b/grammar.js
@@ -188,7 +188,7 @@ module.exports = grammar({
         )
       ),
 
-    shape_key: ($) => choice($.string, $._identifier),
+    shape_key: ($) => choice($.keyword_underscore, $.string, $._identifier),
 
     _identifier: ($) =>
       seq(
@@ -348,6 +348,7 @@ module.exports = grammar({
     _connection_arrowhead_attr_key: ($) =>
       choice("source-arrowhead", "target-arrowhead"),
 
+    keyword_underscore: (_) => "_",
     keyword_classes: (_) => "classes",
     keyword_class: (_) => "class",
     keyword_style: (_) => "style",
diff --git a/queries/highlights.scm b/queries/highlights.scm
index a822220..63383e3 100644
--- a/queries/highlights.scm
+++ b/queries/highlights.scm
@@ -4,12 +4,15 @@
 (shape_key) @variable
 (attr_key) @property
 (reserved) @error
+(class_name) @constant
+
 [
   (keyword_style)
   (keyword_classes)
   (keyword_class)
 ] @keyword
-(class_name) @constant
+
+(keyword_underscore) @keyword.return
 
 ; Literals
 ;-------------------------------------------------------------------------------
diff --git a/src/grammar.json b/src/grammar.json
index 3e39c07..1d2930d 100644
--- a/src/grammar.json
+++ b/src/grammar.json
@@ -853,6 +853,10 @@
     "shape_key": {
       "type": "CHOICE",
       "members": [
+        {
+          "type": "SYMBOL",
+          "name": "keyword_underscore"
+        },
         {
           "type": "SYMBOL",
           "name": "string"
@@ -1818,6 +1822,10 @@
         }
       ]
     },
+    "keyword_underscore": {
+      "type": "STRING",
+      "value": "_"
+    },
     "keyword_classes": {
       "type": "STRING",
       "value": "classes"
@@ -2080,20 +2088,7 @@
       "type": "PATTERN",
       "value": "[0-9]+\\.[0-9]+"
     },
-    "comment": {
-      "type": "CHOICE",
-      "members": [
-        {
-          "type": "SYMBOL",
-          "name": "_line_comment"
-        },
-        {
-          "type": "SYMBOL",
-          "name": "_block_comment"
-        }
-      ]
-    },
-    "_line_comment": {
+    "line_comment": {
       "type": "TOKEN",
       "content": {
         "type": "PREC",
@@ -2150,7 +2145,11 @@
     },
     {
       "type": "SYMBOL",
-      "name": "comment"
+      "name": "line_comment"
+    },
+    {
+      "type": "SYMBOL",
+      "name": "block_comment"
     }
   ],
   "conflicts": [
@@ -2196,7 +2195,7 @@
     },
     {
       "type": "SYMBOL",
-      "name": "_block_comment"
+      "name": "block_comment"
     }
   ],
   "inline": [],
diff --git a/src/node-types.json b/src/node-types.json
index 18d3460..3974383 100644
--- a/src/node-types.json
+++ b/src/node-types.json
@@ -186,6 +186,10 @@
           "type": "escape_sequence",
           "named": true
         },
+        {
+          "type": "keyword_underscore",
+          "named": true
+        },
         {
           "type": "string",
           "named": true
@@ -232,11 +236,6 @@
       ]
     }
   },
-  {
-    "type": "comment",
-    "named": true,
-    "fields": {}
-  },
   {
     "type": "connection",
     "named": true,
@@ -319,6 +318,10 @@
           "type": "escape_sequence",
           "named": true
         },
+        {
+          "type": "keyword_underscore",
+          "named": true
+        },
         {
           "type": "string",
           "named": true
@@ -397,6 +400,10 @@
           "type": "escape_sequence",
           "named": true
         },
+        {
+          "type": "keyword_underscore",
+          "named": true
+        },
         {
           "type": "string",
           "named": true
@@ -513,6 +520,10 @@
     "type": "arrow",
     "named": true
   },
+  {
+    "type": "block_comment",
+    "named": true
+  },
   {
     "type": "bold",
     "named": false
@@ -613,6 +624,10 @@
     "type": "keyword_style",
     "named": true
   },
+  {
+    "type": "keyword_underscore",
+    "named": true
+  },
   {
     "type": "label",
     "named": false
@@ -621,6 +636,10 @@
     "type": "language",
     "named": true
   },
+  {
+    "type": "line_comment",
+    "named": true
+  },
   {
     "type": "link",
     "named": false
diff --git a/src/parser.c b/src/parser.c
index d296682..2a47b89 100644
Binary files a/src/parser.c and b/src/parser.c differ
diff --git a/test/corpus/container.txt b/test/corpus/container.txt
index 81bdf96..5742d7a 100644
--- a/test/corpus/container.txt
+++ b/test/corpus/container.txt
@@ -267,3 +267,25 @@ primty: Primitive types {
     )
   )
 )
+
+================================================================================
+Declare a container with a references to the parent container
+================================================================================
+container: {
+  _.foo: baz
+}
+
+--------------------------------------------------------------------------------
+
+(source_file
+  (container
+    (container_key)
+    (block
+      (container
+        (container_key (keyword_underscore))
+        (dot)
+        (shape (shape_key) (label))
+      )
+    )
+  )
+)
diff --git a/test/corpus/shape.txt b/test/corpus/shape.txt
index 54b87da..4f3dcd5 100644
--- a/test/corpus/shape.txt
+++ b/test/corpus/shape.txt
@@ -278,3 +278,13 @@ bar
   (block_comment)
   (shape (shape_key))
 )
+
+================================================================================
+Declare a shape with a parent container reference
+================================================================================
+_
+--------------------------------------------------------------------------------
+
+(source_file
+  (shape (shape_key (keyword_underscore)))
+)
diff --git a/test/highlight/shape.d2 b/test/highlight/shape.d2
index 06576c7..8468b3e 100644
--- a/test/highlight/shape.d2
+++ b/test/highlight/shape.d2
@@ -33,6 +33,12 @@ foo: Foo Bar {
   bar.baz
   # <- string.special
   #   ^ variable
+
+  _.foo.style.fill: "#8800aa"
+  # <- keyword.return
+  # ^ variable
+  #     ^ keyword
+  #           ^ property
 }
 # <- punctuation.bracket
 
diff --git a/tree-sitter-d2.wasm b/tree-sitter-d2.wasm
index 675efd8..1bff848 100755
Binary files a/tree-sitter-d2.wasm and b/tree-sitter-d2.wasm differ