From dca2abd6cad8976fa154686cd8a7831de679f679 Mon Sep 17 00:00:00 2001
From: Sergey B Kirpichev <skirpichev@gmail.com>
Date: Sun, 21 Jun 2026 11:44:36 +0300
Subject: [PATCH] Cleanup jtheta()

* normalize result
* use extraprec context manager
* add error message
---
 mpmath/functions/theta.py     | 10 +++-------
 mpmath/tests/test_elliptic.py |  6 ++++++
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/mpmath/functions/theta.py b/mpmath/functions/theta.py
index 29252be5..3b349eb9 100644
--- a/mpmath/functions/theta.py
+++ b/mpmath/functions/theta.py
@@ -856,9 +856,7 @@ def jtheta(ctx, n, z, q, derivative=0):
             extra += 2*abs(M)
     cz = 0.5
     extra2 = 50
-    prec0 = ctx.prec
-    try:
-        ctx.prec += extra
+    with ctx.extraprec(extra):
         if n in [1, 2]:
             z_inner = z - ctx.pi/2 if n == 1 else z
             if z.imag:
@@ -882,7 +880,5 @@ def jtheta(ctx, n, z, q, derivative=0):
             else:
                 res = ctx._djacobi_theta3(z, q_inner, nd)
         else:
-            raise ValueError
-    finally:
-        ctx.prec = prec0
-    return res
+            raise ValueError("First argument expected to be 1, 2, 3 or 4")
+    return +res
diff --git a/mpmath/tests/test_elliptic.py b/mpmath/tests/test_elliptic.py
index 83f8dc89..431c01bb 100644
--- a/mpmath/tests/test_elliptic.py
+++ b/mpmath/tests/test_elliptic.py
@@ -213,6 +213,12 @@ def test_jtheta_issue_79():
     r = jtheta(3, 4.5, 0.25, 50)
     assert r.ae('-6148327726309051673317975084654262.4119215720343656')
 
+def test_jtheta_invalid_n():
+    pytest.raises(ValueError, jtheta, 5, 0.5, 0.3)
+    pytest.raises(ValueError, jtheta, 0, 0.5, 0.3)
+    pytest.raises(ValueError, jtheta, 5, 0.5, 0.3)
+    pytest.raises(ValueError, jtheta, 0, 0.5, 0.3)
+
 def test_jtheta_identities():
     """
     Tests the some of the jacobi identidies found in Abramowitz,