|
21 | 21 | Final,
|
22 | 22 | Literal,
|
23 | 23 | cast,
|
24 |
| - overload, |
25 | 24 | )
|
26 | 25 | import warnings
|
27 | 26 |
|
|
85 | 84 | DatetimeArray,
|
86 | 85 | PeriodArray,
|
87 | 86 | )
|
| 87 | +from pandas.core.arrays.datetimes import tz_to_dtype |
88 | 88 | import pandas.core.common as com
|
89 | 89 | from pandas.core.computation.pytables import (
|
90 | 90 | PyTablesExpr,
|
@@ -2170,7 +2170,12 @@ def convert(
|
2170 | 2170 | if "freq" in kwargs:
|
2171 | 2171 | kwargs["freq"] = None
|
2172 | 2172 | new_pd_index = factory(values, **kwargs)
|
2173 |
| - final_pd_index = _set_tz(new_pd_index, self.tz) |
| 2173 | + |
| 2174 | + final_pd_index: Index |
| 2175 | + if self.tz is not None and isinstance(new_pd_index, DatetimeIndex): |
| 2176 | + final_pd_index = new_pd_index.tz_localize("UTC").tz_convert(self.tz) |
| 2177 | + else: |
| 2178 | + final_pd_index = new_pd_index |
2174 | 2179 | return final_pd_index, final_pd_index
|
2175 | 2180 |
|
2176 | 2181 | def take_data(self):
|
@@ -2567,7 +2572,7 @@ def convert(self, values: np.ndarray, nan_rep, encoding: str, errors: str):
|
2567 | 2572 | # reverse converts
|
2568 | 2573 | if dtype.startswith("datetime64"):
|
2569 | 2574 | # recreate with tz if indicated
|
2570 |
| - converted = _set_tz(converted, tz, coerce=True) |
| 2575 | + converted = _set_tz(converted, tz) |
2571 | 2576 |
|
2572 | 2577 | elif dtype == "timedelta64":
|
2573 | 2578 | converted = np.asarray(converted, dtype="m8[ns]")
|
@@ -2948,7 +2953,7 @@ def read_array(self, key: str, start: int | None = None, stop: int | None = None
|
2948 | 2953 | if dtype and dtype.startswith("datetime64"):
|
2949 | 2954 | # reconstruct a timezone if indicated
|
2950 | 2955 | tz = getattr(attrs, "tz", None)
|
2951 |
| - ret = _set_tz(ret, tz, coerce=True) |
| 2956 | + ret = _set_tz(ret, tz) |
2952 | 2957 |
|
2953 | 2958 | elif dtype == "timedelta64":
|
2954 | 2959 | ret = np.asarray(ret, dtype="m8[ns]")
|
@@ -4298,7 +4303,8 @@ def read_column(
|
4298 | 4303 | encoding=self.encoding,
|
4299 | 4304 | errors=self.errors,
|
4300 | 4305 | )
|
4301 |
| - return Series(_set_tz(col_values[1], a.tz), name=column, copy=False) |
| 4306 | + cvs = col_values[1] |
| 4307 | + return Series(cvs, name=column, copy=False) |
4302 | 4308 |
|
4303 | 4309 | raise KeyError(f"column [{column}] not found in the table")
|
4304 | 4310 |
|
@@ -4637,7 +4643,7 @@ def read(
|
4637 | 4643 | if values.ndim == 1 and isinstance(values, np.ndarray):
|
4638 | 4644 | values = values.reshape((1, values.shape[0]))
|
4639 | 4645 |
|
4640 |
| - if isinstance(values, np.ndarray): |
| 4646 | + if isinstance(values, (np.ndarray, DatetimeArray)): |
4641 | 4647 | df = DataFrame(values.T, columns=cols_, index=index_, copy=False)
|
4642 | 4648 | elif isinstance(values, Index):
|
4643 | 4649 | df = DataFrame(values, columns=cols_, index=index_)
|
@@ -4873,54 +4879,21 @@ def _get_tz(tz: tzinfo) -> str | tzinfo:
|
4873 | 4879 | return zone
|
4874 | 4880 |
|
4875 | 4881 |
|
4876 |
| -@overload |
4877 |
| -def _set_tz( |
4878 |
| - values: np.ndarray | Index, tz: str | tzinfo, coerce: bool = False |
4879 |
| -) -> DatetimeIndex: |
4880 |
| - ... |
4881 |
| - |
4882 |
| - |
4883 |
| -@overload |
4884 |
| -def _set_tz(values: np.ndarray | Index, tz: None, coerce: bool = False) -> np.ndarray: |
4885 |
| - ... |
4886 |
| - |
4887 |
| - |
4888 |
| -def _set_tz( |
4889 |
| - values: np.ndarray | Index, tz: str | tzinfo | None, coerce: bool = False |
4890 |
| -) -> np.ndarray | DatetimeIndex: |
| 4882 | +def _set_tz(values: npt.NDArray[np.int64], tz: str | tzinfo | None) -> DatetimeArray: |
4891 | 4883 | """
|
4892 |
| - coerce the values to a DatetimeIndex if tz is set |
4893 |
| - preserve the input shape if possible |
| 4884 | + Coerce the values to a DatetimeArray with appropriate tz. |
4894 | 4885 |
|
4895 | 4886 | Parameters
|
4896 | 4887 | ----------
|
4897 |
| - values : ndarray or Index |
4898 |
| - tz : str or tzinfo |
4899 |
| - coerce : if we do not have a passed timezone, coerce to M8[ns] ndarray |
| 4888 | + values : ndarray[int64] |
| 4889 | + tz : str, tzinfo, or None |
4900 | 4890 | """
|
4901 |
| - if isinstance(values, DatetimeIndex): |
4902 |
| - # If values is tzaware, the tz gets dropped in the values.ravel() |
4903 |
| - # call below (which returns an ndarray). So we are only non-lossy |
4904 |
| - # if `tz` matches `values.tz`. |
4905 |
| - assert values.tz is None or values.tz == tz |
4906 |
| - if values.tz is not None: |
4907 |
| - return values |
4908 |
| - |
4909 |
| - if tz is not None: |
4910 |
| - if isinstance(values, DatetimeIndex): |
4911 |
| - name = values.name |
4912 |
| - else: |
4913 |
| - name = None |
4914 |
| - values = values.ravel() |
4915 |
| - |
4916 |
| - values = DatetimeIndex(values, name=name) |
4917 |
| - values = values.tz_localize("UTC").tz_convert(tz) |
4918 |
| - elif coerce: |
4919 |
| - values = np.asarray(values, dtype="M8[ns]") |
4920 |
| - |
4921 |
| - # error: Incompatible return value type (got "Union[ndarray, Index]", |
4922 |
| - # expected "Union[ndarray, DatetimeIndex]") |
4923 |
| - return values # type: ignore[return-value] |
| 4891 | + assert values.dtype == "i8", values.dtype |
| 4892 | + # Argument "tz" to "tz_to_dtype" has incompatible type "str | tzinfo | None"; |
| 4893 | + # expected "tzinfo" |
| 4894 | + dtype = tz_to_dtype(tz=tz, unit="ns") # type: ignore[arg-type] |
| 4895 | + dta = DatetimeArray._from_sequence(values, dtype=dtype) |
| 4896 | + return dta |
4924 | 4897 |
|
4925 | 4898 |
|
4926 | 4899 | def _convert_index(name: str, index: Index, encoding: str, errors: str) -> IndexCol:
|
|
0 commit comments