Add custom comparison support for dbt-jinja objects

crates/dbt-jinja/minijinja-contrib/src/modules/py_datetime/datetime.rs

@@ -1,10 +1,11 @@
 use std::fmt;
 use std::str::FromStr;
 use std::sync::Arc;
+use std::cmp::Ordering;
 
 use chrono::{DateTime, Datelike, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Timelike, Utc};
 use chrono_tz::Tz;
-use minijinja::{arg_utils::ArgParser, value::Object, Error, ErrorKind, Value};
+use minijinja::{arg_utils::ArgParser, value::{Object, ObjectRepr}, Error, ErrorKind, Value};
 
 use crate::modules::py_datetime::date::PyDate; // your date
 use crate::modules::py_datetime::time::PyTime;
@@ -824,10 +825,24 @@ impl PyDateTime {
 // Implement the `Object` trait for PyDateTime so Jinja can call methods
 //
 impl Object for PyDateTime {
+    fn repr(self: &Arc<Self>) -> ObjectRepr {
+        ObjectRepr::Plain
+    }
+
     fn is_true(self: &Arc<Self>) -> bool {
         true
     }
 
+    fn custom_cmp(self: &Arc<Self>, other: &minijinja::value::DynObject) -> Option<Ordering> {
+        // try to downcast the other object to PyDateTime
+        if let Some(other_dt) = other.downcast_ref::<PyDateTime>() {
+            // compare using timestamps, wrap in Some() to satisfy the trait bounds
+            Some(self.timestamp().total_cmp(&other_dt.timestamp()))
+        } else {
+            None
+        }
+    }
+
     fn call_method(
         self: &Arc<Self>,
         _state: &minijinja::State<'_, '_>,

crates/dbt-jinja/minijinja-contrib/tests/datetime.rs

@@ -285,3 +285,199 @@ fn test_timeformat() {
         .unwrap();
     assert_eq!(expr.eval((), &[]).unwrap().to_string(), "19:37");
 }
+
+#[test]
+fn test_datetime_direct_comparison() {
+    let mut env = minijinja::Environment::new();
+    minijinja_contrib::add_to_environment(&mut env);
+
+    // Test direct comparison operators now work!
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dt1 = modules.datetime.datetime(2025, 1, 1, 12, 0, 0) -%}
+        {%- set dt2 = modules.datetime.datetime(2025, 1, 2, 12, 0, 0) -%}
+        {%- set dt3 = modules.datetime.datetime(2025, 1, 1, 12, 0, 0) -%}
+        dt1 < dt2: {{ dt1 < dt2 }}
+        dt2 > dt1: {{ dt2 > dt1 }}
+        dt1 <= dt2: {{ dt1 <= dt2 }}
+        dt2 >= dt1: {{ dt2 >= dt1 }}
+        dt1 == dt3: {{ dt1 == dt3 }}
+        dt1 != dt2: {{ dt1 != dt2 }}
+        dt1 == dt1: {{ dt1 == dt1 }}
+        "#,
+        &[]
+    ).unwrap();
+
+    let output = tmpl.render(context!{}, &[]).unwrap();
+    assert!(output.contains("dt1 < dt2: true"));
+    assert!(output.contains("dt2 > dt1: true"));
+    assert!(output.contains("dt1 <= dt2: true"));
+    assert!(output.contains("dt2 >= dt1: true"));
+    assert!(output.contains("dt1 == dt3: true"));
+    assert!(output.contains("dt1 != dt2: true"));
+    assert!(output.contains("dt1 == dt1: true"));
+}
+
+#[test]
+fn test_datetime_direct_comparison_timezone() {
+    let mut env = minijinja::Environment::new();
+    minijinja_contrib::add_to_environment(&mut env);
+
+    // Test timezone-aware comparisons with direct operators
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set utc = modules.pytz.timezone('UTC') -%}
+        {%- set eastern = modules.pytz.timezone('US/Eastern') -%}
+        {%- set dt1 = modules.datetime.datetime(2025, 1, 1, 12, 0, 0, tzinfo=utc) -%}
+        {%- set dt2 = modules.datetime.datetime(2025, 1, 1, 7, 0, 0, tzinfo=eastern) -%}
+        {%- set dt3 = modules.datetime.datetime(2025, 1, 1, 8, 0, 0, tzinfo=eastern) -%}
+        same_instant: {{ dt1 == dt2 }}
+        dt1_before_dt3: {{ dt1 < dt3 }}
+        dt3_after_dt1: {{ dt3 > dt1 }}
+        "#,
+        &[]
+    ).unwrap();
+
+    let output = tmpl.render(context!{}, &[]).unwrap();
+    assert!(output.contains("same_instant: true"));
+    assert!(output.contains("dt1_before_dt3: true"));
+    assert!(output.contains("dt3_after_dt1: true"));
+}
+
+#[test]
+fn test_datetime_comparison_with_incompatible_types() {
+    let mut env = minijinja::Environment::new();
+    minijinja_contrib::add_to_environment(&mut env);
+
+    // Test that comparing datetime with non-datetime types returns false
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dt1 = modules.datetime.datetime(2025, 1, 1, 12, 0, 0) -%}
+        {%- set num = 42 -%}
+        {%- set str = "2025-01-01" -%}
+        {%- set none_val = none -%}
+        {%- set date_obj = modules.datetime.date(2025, 1, 1) -%}
+        dt_vs_num: {{ dt1 == num }}
+        dt_vs_str: {{ dt1 == str }}
+        dt_vs_none: {{ dt1 == none_val }}
+        dt_vs_date: {{ dt1 == date_obj }}
+        "#,
+        &[]
+    ).unwrap();
+
+    let output = tmpl.render(context!{}, &[]).unwrap();
+    assert!(output.contains("dt_vs_num: false"));
+    assert!(output.contains("dt_vs_str: false"));
+    assert!(output.contains("dt_vs_none: false"));
+    assert!(output.contains("dt_vs_date: false"));
+}
+
+#[test]
+fn test_datetime_invalid_construction_errors() {
+    let mut env = minijinja::Environment::new();
+    minijinja_contrib::add_to_environment(&mut env);
+
+    // Test invalid month
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dt = modules.datetime.datetime(2025, 13, 1) -%}
+        "#,
+        &[]
+    ).unwrap();
+
+    let result = tmpl.render(context!{}, &[]);
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("Invalid date components"));
+
+    // Test invalid day
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dt = modules.datetime.datetime(2025, 2, 30) -%}
+        "#,
+        &[]
+    ).unwrap();
+
+    let result = tmpl.render(context!{}, &[]);
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("Invalid date components"));
+
+    // Test invalid hour
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dt = modules.datetime.datetime(2025, 1, 1, 25) -%}
+        "#,
+        &[]
+    ).unwrap();
+
+    let result = tmpl.render(context!{}, &[]);
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("Invalid time components"));
+
+    // Test invalid minute
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dt = modules.datetime.datetime(2025, 1, 1, 12, 60) -%}
+        "#,
+        &[]
+    ).unwrap();
+
+    let result = tmpl.render(context!{}, &[]);
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("Invalid time components"));
+}
+
+#[test]
+fn test_datetime_comparison_edge_cases() {
+    let mut env = minijinja::Environment::new();
+    minijinja_contrib::add_to_environment(&mut env);
+
+    // Test edge cases like same datetime, microsecond differences, etc.
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dt1 = modules.datetime.datetime(2025, 1, 1, 12, 0, 0, 0) -%}
+        {%- set dt2 = modules.datetime.datetime(2025, 1, 1, 12, 0, 0, 0) -%}
+        {%- set dt3 = modules.datetime.datetime(2025, 1, 1, 12, 0, 0, 1) -%}
+        same_datetime: {{ dt1 == dt2 }}
+        same_datetime_le: {{ dt1 <= dt2 }}
+        same_datetime_ge: {{ dt1 >= dt2 }}
+        microsecond_diff: {{ dt1 < dt3 }}
+        microsecond_ne: {{ dt1 != dt3 }}
+        "#,
+        &[]
+    ).unwrap();
+
+    let output = tmpl.render(context!{}, &[]).unwrap();
+    assert!(output.contains("same_datetime: true"));
+    assert!(output.contains("same_datetime_le: true"));
+    assert!(output.contains("same_datetime_ge: true"));
+    assert!(output.contains("microsecond_diff: true"));
+    assert!(output.contains("microsecond_ne: true"));
+}
+
+#[test]
+fn test_datetime_sorting() {
+    let mut env = minijinja::Environment::new();
+    minijinja_contrib::add_to_environment(&mut env);
+
+    // Test that datetimes can be sorted correctly
+    let tmpl = env.template_from_str(
+        r#"
+        {%- set dates = [
+            modules.datetime.datetime(2025, 3, 15),
+            modules.datetime.datetime(2025, 1, 1),
+            modules.datetime.datetime(2025, 12, 31),
+            modules.datetime.datetime(2025, 6, 15),
+            modules.datetime.datetime(2025, 1, 1)
+        ] -%}
+        {%- set sorted_dates = dates | sort %}
+        {%- for date in sorted_dates %}
+        {{ date.strftime("%Y-%m-%d") }}
+        {%- endfor %}
+        "#,
+        &[]
+    ).unwrap();
+
+    let output = tmpl.render(context!{}, &[]).unwrap();
+    let dates: Vec<&str> = output.trim().split('\n').map(|s| s.trim()).collect();
+    assert_eq!(dates, vec!["2025-01-01", "2025-01-01", "2025-03-15", "2025-06-15", "2025-12-31"]);
+}

crates/dbt-jinja/minijinja-py/tests/test_basic.py

@@ -2,6 +2,7 @@
 import pytest
 import posixpath
 import types
+from functools import total_ordering
 
 from _pytest.unraisableexception import catch_unraisable_exception
 from minijinja import (
@@ -349,3 +350,27 @@ def test_custom_delimiters():
     )
     rv = env.render_str("<% if true %>${ value }<% endif %><!-- nothing -->", value=42)
     assert rv == "42"
+
+def test_pass_through_sort():
+    @total_ordering
+    class X(object):
+        def __init__(self, value):
+            self.value = value
+
+        def __eq__(self, other):
+            if type(self) is not type(other):
+                return NotImplemented
+            return self.value == other.value
+
+        def __lt__(self, other):
+            if type(self) is not type(other):
+                return NotImplemented
+            return self.value < other.value
+
+        def __str__(self):
+            return str(self.value)
+
+    values = [X(4), X(23), X(42), X(-1)]
+    env = Environment()
+    rv = env.render_str("{{ values|sort|join(',') }}", values=values)
+    assert rv == "-1,4,23,42"

crates/dbt-jinja/minijinja/src/value/mod.rs

@@ -542,6 +542,10 @@ impl PartialEq for Value {
                     if let (Some(a), Some(b)) = (self.as_object(), other.as_object()) {
                         if a.is_same_object(b) {
                             return true;
+                        } else if a.is_same_object_type(b) {
+                            if let Some(rv) = a.custom_cmp(b) {
+                                return rv == Ordering::Equal;
+                            }
                         }
                         match (a.repr(), b.repr()) {
                             (ObjectRepr::Map, ObjectRepr::Map) => {
@@ -629,17 +633,26 @@ impl Ord for Value {
                 Some(ops::CoerceResult::I128(a, b)) => a.cmp(&b),
                 Some(ops::CoerceResult::Str(a, b)) => a.cmp(b),
                 None => {
-                    if let (Some(a), Some(b)) = (self.as_object(), other.as_object()) {
-                        if a.is_same_object(b) {
-                            Ordering::Equal
-                        } else {
-                            match (a.repr(), b.repr()) {
-                                (ObjectRepr::Map, ObjectRepr::Map) => {
-                                    // This is not really correct.  Because the keys can be in arbitrary
-                                    // order this could just sort really weirdly as a result.  However
-                                    // we don't want to pay the cost of actually sorting the keys for
-                                    // ordering so we just accept this for now.
-                                    match (a.try_iter_pairs(), b.try_iter_pairs()) {
+                    // if coerce fails and kinds match, both must be objects
+                    let a = self.as_object().unwrap();
+                    let b = other.as_object().unwrap();
+
+                    if a.is_same_object(b) {
+                        Ordering::Equal
+                    } else {
+                        // if there is a custom comparison, run it.
+                        if a.is_same_object_type(b) {
+                            if let Some(rv) = a.custom_cmp(b) {
+                                return rv;
+                            }
+                        }
+                        match (a.repr(), b.repr()) {
+                            (ObjectRepr::Map, ObjectRepr::Map) => {
+                                // This is not really correct.  Because the keys can be in arbitrary
+                                // order this could just sort really weirdly as a result.  However
+                                // we don't want to pay the cost of actually sorting the keys for
+                                // ordering so we just accept this for now.
+                                match (a.try_iter_pairs(), b.try_iter_pairs()) {
                                         (Some(a), Some(b)) => a.cmp(b),
                                         _ => unreachable!(),
                                     }
@@ -651,11 +664,13 @@ impl Ord for Value {
                                     (Some(a), Some(b)) => a.cmp(b),
                                     _ => unreachable!(),
                                 },
+                                // terrible fallback for plain objects
+                                (ObjectRepr::Plain, ObjectRepr::Plain) => {
+                                    a.to_string().cmp(&b.to_string())
+                                }
+                                // should not happen
                                 (_, _) => unreachable!(),
-                            }
                         }
-                    } else {
-                        unreachable!()
                     }
                 }
             },