From f04adfc661a3e5f6a1d75e1f5035686652e6f337 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Wed, 17 Jan 2024 14:50:39 -0500
Subject: [PATCH] Use type definitions for type checks

---
 src/abstract_tree/assignment.rs | 11 +++++++----
 tests/types.rs                  |  2 --
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/abstract_tree/assignment.rs b/src/abstract_tree/assignment.rs
index 2daa5b2..85fd854 100644
--- a/src/abstract_tree/assignment.rs
+++ b/src/abstract_tree/assignment.rs
@@ -38,10 +38,13 @@ impl AbstractTree for Assignment {
         let statement = Statement::from_syntax(statement_node, source, context)?;
 
         if let AssignmentOperator::Equal = operator {
-            context.set_type(
-                identifier.inner().clone(),
-                statement.expected_type(context)?,
-            )?;
+            let r#type = if let Some(definition) = &type_definition {
+                definition.inner().clone()
+            } else {
+                statement.expected_type(context)?
+            };
+
+            context.set_type(identifier.inner().clone(), r#type)?;
         }
 
         Ok(Assignment {
diff --git a/tests/types.rs b/tests/types.rs
index b1e88a6..8dabad4 100644
--- a/tests/types.rs
+++ b/tests/types.rs
@@ -12,8 +12,6 @@ fn simple_type_check() {
 
 #[test]
 fn argument_count_check() {
-    env_logger::builder().is_test(true).try_init().unwrap();
-
     let source = "
             foo = (x <int>) <int> {
                 x