Skip to content

TYP: Fixed & improved numpy.dtype.__new__ - #27139

Merged
charris merged 7 commits into
numpy:mainfrom
jorenham:typing/ufunc-constructor-specialization
Aug 30, 2024
Merged

TYP: Fixed & improved numpy.dtype.__new__#27139
charris merged 7 commits into
numpy:mainfrom
jorenham:typing/ufunc-constructor-specialization

Conversation

@jorenham

@jorenham jorenham commented Aug 8, 2024 •
edited
Loading

Copy link
Copy Markdown
Member

The changes are mostly limited to the numpy.dtype.__new__ overloads:

  • Support for StringDType: Type-checkers will now allow dtype(np.dtypes.StringDType) (which previously would've been rejected), and correctly infer the return type of e.g. dtype('T') as -> StringDType.
  • Full support for input that's a union of char-codes, ctypes, or python builtin types, analogous to the scalar-type hierarchy. So calling dtype(arg) with some arg: str | bytes now returns a dtype[character] (previously this was dtype[object_]). And passing arg: Literal["int", "c16"] now results in dtype[number[Any]] (which was dtype[Any]).
  • Fix dtype(memoryview), which was previously incorrectly annotated as dtype[object_], instead of dtype[void]. This requires type-checkers to be PEP 688 compliant, otherwise these overloads are skipped, and the inferred type will be same same as it would've been before.
    See also Drop bytes/bytearray and bytes/memoryview promotions (PEP 688) python/mypy#15313 .
  • The overlapping overloads issue has been fixed, as explained in TYP: Overlapping overloads with incompatible return types #27032. The inevitable consequence of this, is that dtype(float) now returns dtype[float64 | int_ | np.bool]. This is necessary because of the "special" way in which float and complex should be treated, according to the typing docs.
  • This also fixes many cases where the return type was incorrectly inferred as dtype[object_] instead of dtype[Any]. This was caused by the last overload with signature (type[object], ...) -> dtype[object_], which also matches e.g. int | str.

This additionally adds several new literals to numpy._typing._char_codes for the abstract scalar types, and for StringDType.

A quick example of why this is needed:

# pyright: reportInvalidStubStatement=false
import numpy as np
from numpy import dtype, dtypes


"""Constructing a `StringDType` through `numpy.dtype`"""
reveal_type(dtype("T"))
# PRIOR: Type of "dtype("T")" is "dtype[Any]"
# AFTER: Type of "dtype("T")" is "StringDType"
reveal_type(dtype(dtypes.StringDType()))
# PRIOR: Type of "dtype(dtypes.StringDType())" is "dtype[Unknown]"
# AFTER: Type of "dtype(dtypes.StringDType())" is "StringDType"
reveal_type(dtype(dtypes.BoolDType()))  # `numpy.dtypes` isn't the issue
# PRIOR: Type of "dtype(dtypes.BoolDType())" is "dtype[bool]"
# AFTER: Type of "dtype(dtypes.BoolDType())" is "dtype[bool]"


"""Char-code unions."""
from typing import Literal
char_SUV: Literal["S"] | Literal["U"] | Literal["V"]
reveal_type(dtype(char_SUV))
# PRIOR: Type of "dtype(char_SUV)" is "dtype[Any]"
# AFTER: Type of "dtype(char_SUV)" is "dtype[flexible]"

"""Void from `memoryview`"""
reveal_type(dtype(memoryview))
# PRIOR: Type of "dtype(memoryview)" is "dtype[object_]"
# AFTER: Type of "dtype(memoryview)" is "dtype[void]"

"""Overlapping overloads with incompatible return types"""
from random import choice
reveal_type(dtype(choice([np.bool, np.float64])))
# PRIOR: Type of "dtype(choice([np.bool, np.float64]))" is "dtype[bool | floating[_64Bit]]"
# AFTER: Type of "dtype(choice([np.bool, np.float64]))" is "dtype[bool | floating[_64Bit]]"
reveal_type(dtype(choice([bool, int, float])))
# PRIOR: Type of "dtype(choice([bool, int, float]))" is "dtype[floating[_64Bit]]"
# AFTER: Type of "dtype(choice([bool, int, float]))" is "dtype[floating[_64Bit] | signedinteger[Any] | bool]"

"""Union of scalar types"""
reveal_type(dtype(choice([np.bytes_, np.str_])))
# PRIOR: Type of "dtype(choice([np.bytes_, np.str_]))" is "dtype[bytes_ | str_]"
# AFTER: Type of "dtype(choice([np.bytes_, np.str_]))" is "dtype[bytes_ | str_]"
reveal_type(dtype(choice([bytes, str])))
# PRIOR: Type of "dtype(choice([bytes, str]))" is "dtype[object_]"
# AFTER: Type of "dtype(choice([bytes, str]))" is "dtype[character]"
reveal_type(dtype(choice([bool, np.bool])))
# PRIOR: Type of "dtype(choice([bool, np.bool]))" is "dtype[object_]"
# AFTER: Type of "dtype(choice([bool, np.bool]))" is "dtype[bool]"

"""Optional scalar types"""
reveal_type(dtype(choice([None, float])))
# PRIOR: Type of "dtype(choice([None, float]))" is "dtype[floating[_64Bit]]"
# AFTER: Type of "dtype(choice([None, float]))" is "dtype[floating[_64Bit] | signedinteger[Any] | bool]"
reveal_type(dtype(choice([None, np.float64])))
# PRIOR: Type of "dtype(choice([None, np.float64]))" is "dtype[Unknown]"
# AFTER: Type of "dtype(choice([None, np.float64]))" is "dtype[floating[_64Bit]]"

This .pyi was ran through basedpyright 1.15.2 (pyright 1.1.375), and uses the current main branch as the "prior".

@jorenham
jorenham requested a review from seberg August 9, 2024 18:38
@jorenham
jorenham force-pushed the typing/ufunc-constructor-specialization branch from 746e0b1 to d816ff5 Compare August 12, 2024 21:55
@jorenham
jorenham force-pushed the typing/ufunc-constructor-specialization branch from d816ff5 to 9f03471 Compare August 15, 2024 10:45
@charris
charris merged commit af75b6f into numpy:main Aug 30, 2024
@charris

charris commented Aug 30, 2024

Copy link
Copy Markdown
Member

Thanks @jorenham .

@jorenham
jorenham deleted the typing/ufunc-constructor-specialization branch August 30, 2024 19:46
@jakevdp jakevdp mentioned this pull request Dec 18, 2024
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@seberg seberg Awaiting requested review from seberg

Assignees

No one assigned

Labels

41 - Static typing

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jorenham @charris