compute retrievability for all states (#61)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fsrs"
-version = "3.0.0"
+version = "3.1.0"
 description = "Free Spaced Repetition Scheduler"
 readme = "README.md"
 authors = [{ name = "Jarrett Ye", email = "[email protected]" }]

src/fsrs/models.py

@@ -262,24 +262,27 @@ def from_dict(source_dict: dict[str, Any]) -> "Card":
             last_review,
         )
 
-    def get_retrievability(self, now: datetime) -> Optional[float]:
+    def get_retrievability(self, now: Optional[datetime] = None) -> float:
         """
         Calculates the Card object's current retrievability for a given date and time.
 
         Args:
             now (datetime): The current date and time
 
         Returns:
-            Optional[float]: The retrievability of the Card object if it's in the Review state, otherwise, will return None.
+            float: The retrievability of the Card object.
         """
         DECAY = -0.5
         FACTOR = 0.9 ** (1 / DECAY) - 1
 
-        if self.state == State.Review:
+        if now is None:
+            now = datetime.now(timezone.utc)
+
+        if self.state in (State.Learning, State.Review, State.Relearning):
             elapsed_days = max(0, (now - self.last_review).days)
             return (1 + FACTOR * elapsed_days / self.stability) ** DECAY
         else:
-            return None
+            return 0
 
 
 @dataclass

tests/test_fsrs.py

@@ -1,4 +1,4 @@
-from fsrs import *
+from fsrs import FSRS, Card, ReviewLog, State, Rating
 from datetime import datetime, timedelta, timezone
 import json
 import pytest
@@ -310,3 +310,31 @@ def test_custom_scheduler_args(self):
         assert f2.p.w == w2
         assert f2.p.request_retention == request_retention2
         assert f2.p.maximum_interval == maximum_interval2
+
+    def test_retrievability(self):
+        f = FSRS()
+
+        card = Card()
+
+        # retrievabiliy of New card
+        assert card.state == State.New
+        retrievability = card.get_retrievability()
+        assert retrievability == 0
+
+        # retrievabiliy of Learning card
+        card, _ = f.review_card(card, Rating.Good)
+        assert card.state == State.Learning
+        retrievability = card.get_retrievability()
+        assert 0 <= retrievability <= 1
+
+        # retrievabiliy of Review card
+        card, _ = f.review_card(card, Rating.Good)
+        assert card.state == State.Review
+        retrievability = card.get_retrievability()
+        assert 0 <= retrievability <= 1
+
+        # retrievabiliy of Relearning card
+        card, _ = f.review_card(card, Rating.Again)
+        assert card.state == State.Relearning
+        retrievability = card.get_retrievability()
+        assert 0 <= retrievability <= 1