gmpy.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104
  1. import os
  2. from typing import Tuple as tTuple, Type
  3. import mpmath.libmp as mlib
  4. from sympy.external import import_module
  5. __all__ = [
  6. # GROUND_TYPES is either 'gmpy' or 'python' depending on which is used. If
  7. # gmpy is installed then it will be used unless the environment variable
  8. # SYMPY_GROUND_TYPES is set to something other than 'auto', 'gmpy', or
  9. # 'gmpy2'.
  10. 'GROUND_TYPES',
  11. # If HAS_GMPY is 0, no supported version of gmpy is available. Otherwise,
  12. # HAS_GMPY will be 2 for gmpy2 if GROUND_TYPES is 'gmpy'. It used to be
  13. # possible for HAS_GMPY to be 1 for gmpy but gmpy is no longer supported.
  14. 'HAS_GMPY',
  15. # SYMPY_INTS is a tuple containing the base types for valid integer types.
  16. # This is either (int,) or (int, type(mpz(0))) depending on GROUND_TYPES.
  17. 'SYMPY_INTS',
  18. # MPQ is either gmpy.mpq or the Python equivalent from
  19. # sympy.external.pythonmpq
  20. 'MPQ',
  21. # MPZ is either gmpy.mpz or int.
  22. 'MPZ',
  23. # Either the gmpy or the mpmath function
  24. 'factorial',
  25. # isqrt from gmpy or mpmath
  26. 'sqrt',
  27. ]
  28. #
  29. # SYMPY_GROUND_TYPES can be gmpy, gmpy2, python or auto
  30. #
  31. GROUND_TYPES = os.environ.get('SYMPY_GROUND_TYPES', 'auto').lower()
  32. #
  33. # Try to import gmpy2 by default. If gmpy or gmpy2 is specified in
  34. # SYMPY_GROUND_TYPES then warn if gmpy2 is not found. In all cases there is a
  35. # fallback based on pure Python int and PythonMPQ that should still work fine.
  36. #
  37. if GROUND_TYPES in ('auto', 'gmpy', 'gmpy2'):
  38. # Actually import gmpy2
  39. gmpy = import_module('gmpy2', min_module_version='2.0.0',
  40. module_version_attr='version', module_version_attr_call_args=())
  41. # Warn if user explicitly asked for gmpy but it isn't available.
  42. if gmpy is None and GROUND_TYPES in ('gmpy', 'gmpy2'):
  43. from warnings import warn
  44. warn("gmpy library is not installed, switching to 'python' ground types")
  45. elif GROUND_TYPES == 'python':
  46. # The user asked for Python so ignore gmpy2 module.
  47. gmpy = None
  48. else:
  49. # Invalid value for SYMPY_GROUND_TYPES. Ignore the gmpy2 module.
  50. from warnings import warn
  51. warn("SYMPY_GROUND_TYPES environment variable unrecognised. "
  52. "Should be 'python', 'auto', 'gmpy', or 'gmpy2'")
  53. gmpy = None
  54. #
  55. # At this point gmpy will be None if gmpy2 was not successfully imported or if
  56. # the environment variable SYMPY_GROUND_TYPES was set to 'python' (or some
  57. # unrecognised value). The two blocks below define the values exported by this
  58. # module in each case.
  59. #
  60. SYMPY_INTS: tTuple[Type, ...]
  61. if gmpy is not None:
  62. HAS_GMPY = 2
  63. GROUND_TYPES = 'gmpy'
  64. SYMPY_INTS = (int, type(gmpy.mpz(0)))
  65. MPZ = gmpy.mpz
  66. MPQ = gmpy.mpq
  67. factorial = gmpy.fac
  68. sqrt = gmpy.isqrt
  69. else:
  70. from .pythonmpq import PythonMPQ
  71. HAS_GMPY = 0
  72. GROUND_TYPES = 'python'
  73. SYMPY_INTS = (int,)
  74. MPZ = int
  75. MPQ = PythonMPQ
  76. factorial = lambda x: int(mlib.ifac(x))
  77. sqrt = lambda x: int(mlib.isqrt(x))