Diff of /edsnlp/extensions.py [000000] .. [cad161]

Switch to side-by-side view

--- a
+++ b/edsnlp/extensions.py
@@ -0,0 +1,45 @@
+import warnings
+from datetime import date, datetime
+
+from dateutil.parser import parse as parse_date
+from spacy.tokens import Doc
+
+if not Doc.has_extension("note_id"):
+    Doc.set_extension("note_id", default=None)
+
+
+def set_note_datetime(doc, dt):
+    try:
+        if type(dt) is datetime:
+            pass
+        elif isinstance(dt, str):
+            dt = parse_date(dt)
+        elif isinstance(dt, (int, float)):
+            dt = datetime.fromtimestamp(dt)
+        elif isinstance(dt, date):
+            dt = datetime(dt.year, dt.month, dt.day)
+        elif dt is None:
+            pass
+        key = doc._._get_key("note_datetime")
+        doc.doc.user_data[key] = dt
+        return
+    except Exception:
+        pass
+
+    warnings.warn(f"Cannot cast {dt} as a note datetime", UserWarning)
+
+
+def get_note_datetime(doc):
+    key = doc._._get_key("note_datetime")
+    return doc.user_data.get(key, None)
+
+
+if not Doc.has_extension("note_datetime"):
+    Doc.set_extension(
+        "note_datetime",
+        getter=get_note_datetime,
+        setter=set_note_datetime,
+    )
+
+if not Doc.has_extension("birth_datetime"):
+    Doc.set_extension("birth_datetime", default=None)