mspu 0.0.1__py3-none-any.whl

mspu/__init__.py

@@ -0,0 +1,6 @@
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version('mspu')
+except PackageNotFoundError:
+    __version__ = '0.0.0'  # Fallback for development

mspu/data/__init__.py

@@ -0,0 +1,5 @@
+from .rand_df import gen_rand_df
+
+__all__ = [
+    'gen_rand_df',
+]

mspu/data/rand_df.py

@@ -0,0 +1,346 @@
+import string
+import numpy as np
+import pandas as pd
+from collections.abc import Iterable
+
+
+def gen_rand_strs(
+    rng: np.random.Generator,
+    str_cnt: int,
+    str_len: tuple[int, int],
+    str_chars: list[str],
+) -> list[str]:
+    str_lens = rng.integers(
+        low=str_len[0], high=str_len[1], size=str_cnt, endpoint=True
+    )
+    rand_strs = [''.join(rng.choice(str_chars, size=str_len)) for str_len in str_lens]
+    return rand_strs
+
+
+def gen_str_vals(
+    size: int,
+    rng: np.random.Generator,
+    str_cnt: int = None,
+    str_len: int | tuple[int, int] = None,
+    str_chars: list[str] = None,
+    col_strs: list[str] = None,
+) -> np.ndarray:
+    if str_cnt is None:
+        str_cnt = 10
+    if str_len is None:
+        str_len = (5, 5)
+    elif isinstance(str_len, int):
+        str_len = (str_len, str_len)
+    if col_strs is None:
+        if str_chars is None:
+            str_chars = [c for c in string.ascii_letters + string.digits]
+        col_strs = gen_rand_strs(rng, str_cnt, str_len, str_chars)
+    val = rng.choice(col_strs, size=size)
+    return val
+
+
+def gen_ts_vals(
+    size: int,
+    rng: np.random.Generator,
+    start_date: str = None,
+    end_date: str = None,
+    freq: str = None,
+    random: bool = False,
+) -> np.ndarray:
+    if start_date is None:
+        start_date = '2024-01-01'
+    if end_date is None:
+        end_date = '2025-01-01'
+    if freq is None:
+        freq = 'D'
+    if random is None:
+        random = False
+    val = pd.date_range(start_date, end_date, freq=freq, inclusive='left')[:size]
+    if random:
+        val = rng.choice(val, size=size)
+    return val
+
+
+def gen_num_vals(
+    size: int,
+    rng: np.random.Generator,
+    low: int | float = None,
+    high: int | float = None,
+    dtype: str = None,
+) -> np.ndarray:
+    if low is None:
+        low = 0
+    if high is None:
+        high = 2
+    func = rng.integers if dtype[0] == 'i' else rng.uniform
+    vals = func(low=low, high=high, size=size)
+    return vals
+
+
+def gen_missing_vals(
+    vals: np.ndarray,
+    rng: np.random._generator.Generator,
+    dtype: str,
+    missing_pct: float = None,
+) -> np.ndarray:
+    if missing_pct is None or missing_pct <= 0 or missing_pct >= 1:
+        return vals
+    if dtype == 's':
+        missing_val = None
+    elif dtype == 't':
+        missing_val = np.datetime64('NaT')
+    else:
+        missing_val = np.nan
+    if dtype == 'i':
+        vals = vals.astype(np.float64)
+    mask = rng.uniform(size=len(vals)) <= missing_pct
+    vals[mask] = missing_val
+    return vals
+
+
+def sanitize_parameters(
+    name_prefix: str,
+    params: dict,
+    par_names: list[str],
+):
+    if isinstance(params, int):
+        cnt = params
+        par = {}
+    else:
+        cnt = params['count']
+        par = {
+            k: (
+                list(v) + [None] * (cnt - len(v))
+                if isinstance(v, Iterable) and not isinstance(v, str)
+                else [v] * cnt
+            )
+            for k, v in params.items()
+        }
+    if name_prefix in ('i', 'f'):
+        par_names += ['dtype']
+        par['dtype'] = [name_prefix] * cnt
+    default_val = [None] * cnt
+    parameters = [
+        {key: par.get(key, default_val)[i] for key in par_names} for i in range(cnt)
+    ]
+    col_names = par.get('name')
+    if col_names is None or col_names.count(None) > 1 or col_names.count('') > 1:
+        col_names = [f'{name_prefix}{i}' for i in range(1, cnt + 1)]
+    col_missing_pcts = par.get('missing_pct', default_val)
+    return col_names, parameters, col_missing_pcts
+
+
+def gen_rand_df(
+    nrow: int,
+    str_cols: dict = None,
+    ts_cols: dict = None,
+    int_cols: dict = None,
+    float_cols: dict = None,
+    rand_seed: int = 11,
+) -> pd.DataFrame:
+    """
+    Generate a random DataFrame with specified column types and parameters.
+
+    Parameters
+    ----------
+    nrow : int
+        Number of rows in the generated DataFrame.
+    str_cols : dict or int, optional
+        Parameters for string columns. If an integer is provided, it specifies the
+        number of string columns to generate with default parameters. If a dictionary
+        is provided, it should include the following keys:
+
+        - count: number of columns
+        - name: list of column names
+        - str_len: int or tuple of (min, max) string length
+        - str_chars: list of characters to use in strings
+        - col_strs: list of lists of strings to sample from for each column (if
+          provided, overrides str_len and str_chars)
+    ts_cols : dict or int, optional
+        Parameters for timestamp columns. If an integer is provided, it specifies the
+        number of timestamp columns to generate with default parameters. If a dictionary
+        is provided, it should include the following keys:
+
+        - count: number of columns
+        - name: list of column names
+        - start_date: start date for timestamp generation
+        - end_date: end date for timestamp generation
+        - freq: frequency for timestamp generation (e.g. 'D' for daily)
+        - random: whether to sample timestamps randomly from the generated range
+    int_cols : dict or int, optional
+        Parameters for integer columns. If an integer is provided, it specifies the
+        number of integer columns to generate with default parameters. If a dictionary
+        is provided, it should include the following keys:
+
+        - count: number of columns
+        - name: list of column names
+        - low: lower bound for random number generation
+        - high: upper bound for random number generation
+        - missing_pct: percentage of values to set as missing (NaN or None)
+    float_cols : dict or int, optional
+        Parameters for float columns. If an integer is provided, it specifies the
+        number of float columns to generate with default parameters. If a dictionary
+        is provided, it should include the following keys:
+
+        - count: number of columns
+        - name: list of column names
+        - low: lower bound for random number generation
+        - high: upper bound for random number generation
+        - missing_pct: percentage of values to set as missing (NaN or None)
+    rand_seed : int, optional
+        Random seed for reproducibility (default is 11).
+
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame with the specified random data.
+
+    Examples
+    --------
+
+    Simply specify the number of columns for each type.
+
+    >>> import pandas as pd
+    >>> from mspu.data import gen_rand_df
+    >>> from mspu.pandas import pd_ht
+    >>> pd.DataFrame.ht = pd_ht
+    >>> df = gen_rand_df(
+    ...     nrow=365 * 24 * 60,
+    ...     str_cols=1,
+    ...     ts_cols=2,
+    ...     int_cols=1,
+    ...     float_cols=2,
+    ... )
+    >>> df.ht()
+    shape: (525600, 6)
+            s1            t1         t2  i1        f1        f2
+    0       IZD8v 2024-01-01 2024-01-01   0  1.593760  1.648840
+    1       P9r1i 2024-01-02 2024-01-02   0  1.622772  1.943180
+    525598  iU68M        NaT        NaT   1  1.277124  0.093818
+    525599  P9r1i        NaT        NaT   1  1.248407  0.601518
+
+    Provide detailed parameters for each column type.
+
+    >>> import pandas as pd
+    >>> from mspu.data import gen_rand_df
+    >>> from mspu.pandas import pd_ht
+    >>> pd.DataFrame.ht = pd_ht
+    >>> d2 = gen_rand_df(
+    ...     nrow=10,
+    ...     str_cols={
+    ...         'count': 2,
+    ...         'name': ['country', 'color'],
+    ...         'str_len': [3, (3, 9)],
+    ...         'str_cnt': [2, 5],
+    ...         'col_strs': [['UK', 'US', 'AU'], ['blue', 'black', 'red']],
+    ...     },
+    ...     ts_cols={
+    ...         'count': 2,
+    ...         'name': ['start_date', 'end_date'],
+    ...         'start_date': ['2020-01-01', '2024-01-01'],
+    ...         'end_date': ['2021-01-01', '2025-01-01'],
+    ...         'freq': 'QS',
+    ...         'random': False,
+    ...     },
+    ...     int_cols={
+    ...         'count': 1,
+    ...         'name': ['quantity'],
+    ...         'low': [0],
+    ...         'high': [100],
+    ...         'missing_pct': [0.3],
+    ...     },
+    ...     float_cols={
+    ...         'count': 2,
+    ...         'name': ['price', 'charge'],
+    ...         'low': [1, 0.1],
+    ...         'high': [100, 0.9],
+    ...         'missing_pct': [0.3, 0.2],
+    ...     },
+    ... )
+    >>> d2.ht(1)
+    shape: (10, 7)
+      country  color start_date   end_date  quantity      price    charge
+    0      UK  black 2020-01-01 2024-01-01      86.0        NaN  0.657889
+    9      UK  black        NaT        NaT      13.0  43.564921  0.776712
+    """
+    col_types = ['s', 't', 'i', 'f']
+    col_params = [str_cols, ts_cols, int_cols, float_cols]
+    funcs = [gen_str_vals, gen_ts_vals, gen_num_vals, gen_num_vals]
+    par_names = [
+        ['str_cnt', 'str_len', 'str_chars', 'col_strs'],
+        ['start_date', 'end_date', 'freq', 'random'],
+        ['low', 'high'],
+        ['low', 'high'],
+    ]
+    dfs = []
+    rand_rng = np.random.default_rng(seed=rand_seed)
+    for i, params in enumerate(col_params):
+        if params is None:
+            continue
+        col_names, col_params, col_missing_pcts = sanitize_parameters(
+            col_types[i], params, par_names[i]
+        )
+        df = pd.DataFrame(
+            {
+                col: gen_missing_vals(
+                    funcs[i](nrow, rand_rng, **col_params[j]),
+                    rand_rng,
+                    col_types[i],
+                    col_missing_pcts[j],
+                )
+                for j, col in enumerate(col_names)
+            }
+        )
+        dfs.append(df)
+    if len(dfs) == 0:
+        df = pd.DataFrame()
+    else:
+        df = pd.concat(dfs, axis=1)
+    return df
+
+
+if __name__ == '__main__':
+    # example 1
+    df = gen_rand_df(
+        nrow=365 * 24 * 60,
+        str_cols=1,
+        ts_cols=2,
+        int_cols=1,
+        float_cols=2,
+    )
+    print(df[:2])
+
+    # example 2
+    d2 = gen_rand_df(
+        nrow=10,
+        str_cols={
+            'count': 2,
+            'name': ['country', 'color'],
+            'str_len': [3, (3, 9)],
+            'str_cnt': [2, 5],
+            'col_strs': [['UK', 'US', 'AU'], ['blue', 'black', 'red']],
+        },
+        ts_cols={
+            'count': 2,
+            'name': ['start_date', 'end_date'],
+            'start_date': ['2020-01-01', '2024-01-01'],
+            'end_date': ['2021-01-01', '2025-01-01'],
+            'freq': 'QS',
+            'random': False,
+        },
+        int_cols={
+            'count': 1,
+            'name': ['quantity'],
+            'low': [0],
+            'high': [100],
+            'missing_pct': [0.3],
+        },
+        float_cols={
+            'count': 2,
+            'name': ['price', 'charge'],
+            'low': [1, 0.1],
+            'high': [100, 0.9],
+            'missing_pct': [0.3, 0.2],
+        },
+    )
+    print(d2[:3])

mspu/pandas/__init__.py

@@ -0,0 +1,10 @@
+from .datetime import explode_date_range
+from .parquet import pa_mod
+from .utils import pd_ht, df_diffs
+
+__all__ = [
+    'pd_ht',
+    'df_diffs',
+    'explode_date_range',
+    'pa_mod',
+]

mspu/pandas/datetime.py

@@ -0,0 +1,262 @@
+import pandas as pd
+from typing import Literal
+
+
+def explode_date_range(
+    df: pd.DataFrame,
+    start_date_col: str,
+    end_date_col: str,
+    date_col: str = 'ts',
+    freq: str = '30min',
+    start_date_offset: pd.DateOffset = None,
+    end_date_offset: pd.DateOffset = None,
+    start_date_roll: Literal['backward', 'forward'] | None = None,
+    end_date_roll: Literal['backward', 'forward'] | None = None,
+    min_date: str | pd.Timestamp = None,
+    max_date: str | pd.Timestamp = None,
+    inclusive: Literal['both', 'left', 'right', 'neither'] = 'both',
+    drop_index: bool = True,
+    drop_date_cols: bool = True,
+) -> pd.DataFrame:
+    """
+    Explode DataFrame start/end date columns to date column.
+
+    Parameters
+    ----------
+    df:
+        The input DataFrame with start/end date columns/index levels.
+    start_date_col:
+        The column name in the DataFrame for the start date.
+    end_date_col:
+        The column name in the DataFrame for the end date.
+    date_col:
+        The column name in the DataFrame for the new date.
+    freq:
+        The frequency of the new date column.
+    start_date_offset:
+        The date offset for the start date column.
+    end_date_offset:
+        The date offset for the end date column.
+    start_date_roll:
+        Roll the start_date to the start of the current/next period.
+    end_date_roll:
+        Roll the end date to the start of the current/next period.
+    min_date:
+        The min value of the start date after offset.
+    max_date:
+        The max value of the end date after offset.
+    inclusive:
+        Include boundaries; Whether to set each bound as closed or open.
+    drop_index:
+        This flag should be False when the input DataFrame has meaningful index.
+    drop_date_cols:
+        Whether to drop the start_date_col and end_date_col or not.
+
+    Returns
+    -------
+    pl.DataFrame
+        The DataFrame same as the input but with
+        start/end date columns replaced by the new date column
+
+    Examples
+    --------
+    Simple usage.
+
+    >>> df = pd.DataFrame({
+    ...     'start_date': ['2023-01-01', '2023-01-02'],
+    ...     'end_date': ['2023-01-01 01:00', '2023-01-02 02:00'],
+    ... })
+    >>> df_exploded = explode_date_range(df, 'start_date', 'end_date', freq='1h')
+    >>> print(df_exploded)
+                       ts
+    0 2023-01-01 00:00:00
+    1 2023-01-01 01:00:00
+    2 2023-01-02 00:00:00
+    3 2023-01-02 01:00:00
+    4 2023-01-02 02:00:00
+
+    With offset for start and end date.
+
+    >>> df = pd.DataFrame({
+    ...     'start_date': ['2023-01-01', '2023-01-02'],
+    ...     'end_date': ['2023-01-01 01:00:00', '2023-01-02 02:00:00'],
+    ... })
+    >>> df_exploded = explode_date_range(
+    ...     df=df,
+    ...     start_date_col='start_date',
+    ...     end_date_col='end_date',
+    ...     freq='1h',
+    ...     start_date_offset=pd.DateOffset(hours=1),
+    ...     end_date_offset=pd.DateOffset(hours=-1),
+    ... )
+    >>> print(df_exploded)
+                       ts
+    0 2023-01-02 01:00:00
+
+    With rolling start and end date to the nearest frequency (hour).
+
+    >>> df = pd.DataFrame({
+    ...     'start_date': ['2023-01-01 00:30', '2023-01-02 00:30'],
+    ...     'end_date': ['2023-01-01 01:30', '2023-01-02 02:30'],
+    ... })
+    >>> df_exploded = explode_date_range(
+    ...     df=df,
+    ...     start_date_col='start_date',
+    ...     end_date_col='end_date',
+    ...     freq='1h',
+    ...     start_date_roll='forward',
+    ...     end_date_roll='backward',
+    ... )
+    >>> print(df_exploded)
+                       ts
+    0 2023-01-01 01:00:00
+    1 2023-01-02 01:00:00
+    2 2023-01-02 02:00:00
+    """
+    if not drop_index:
+        levels_old = list(df.index.names)
+        index_names = [
+            f'_idx{i}' if name is None else name for i, name in enumerate(levels_old)
+        ]
+        df = df.rename_axis(index_names, axis=0)
+        if start_date_col in df.columns and end_date_col in df.columns:
+            levels = index_names
+        else:
+            if drop_date_cols:
+                levels = [
+                    level
+                    for level in index_names
+                    if level not in (start_date_col, end_date_col)
+                ]
+                levels_old = [
+                    level
+                    for level in levels_old
+                    if level not in (start_date_col, end_date_col)
+                ]
+            else:
+                levels = index_names
+                # move start/end_date_col to index if one exists
+                if start_date_col in df.columns:
+                    levels += [start_date_col]
+                    levels_old += [start_date_col]
+                elif end_date_col in df.columns:
+                    levels += [end_date_col]
+                    levels_old += [end_date_col]
+            levels += [date_col]
+            levels_old += [date_col]
+
+    df = df.reset_index(drop=drop_index).astype(
+        {
+            start_date_col: 'datetime64[ns]',
+            end_date_col: 'datetime64[ns]',
+        }
+    )
+
+    # offset start date
+    if start_date_offset is not None:
+        df[start_date_col] += start_date_offset
+
+    # roll start date
+    if start_date_roll is not None:
+        roll_freq = freq if freq[-1] != 'S' else freq[:-1]
+        extra_period = 0 if start_date_roll == 'backward' else 1
+        df[start_date_col] = (
+            df[start_date_col].dt.to_period(roll_freq) + extra_period
+        ).dt.start_time
+
+    # limit start_date and replace null with min_date
+    if min_date is not None:
+        min_date = pd.to_datetime(min_date, dayfirst=True)
+        df[start_date_col] = (
+            df[start_date_col]
+            .fillna(min_date)
+            .where(df[start_date_col] > min_date, min_date)
+        )
+
+    # offset end date
+    if end_date_offset is not None:
+        df[end_date_col] += end_date_offset
+
+    # roll end date
+    if end_date_roll is not None:
+        roll_freq = freq if freq[-1] != 'S' else freq[:-1]
+        extra_period = 0 if end_date_roll == 'backward' else 1
+        df[end_date_col] = (
+            df[end_date_col].dt.to_period(roll_freq) + extra_period
+        ).dt.start_time
+
+    # limit end_date and replace null with max_date
+    if max_date is not None:
+        max_date = pd.to_datetime(max_date, dayfirst=True)
+        df[end_date_col] = (
+            df[end_date_col]
+            .fillna(max_date)
+            .where(df[end_date_col] < max_date, max_date)
+        )
+
+    # FIXME: special regarding pandas version
+    pd_version = tuple(int(n) for n in pd.__version__.split('.')[:3])
+    if pd_version >= (1, 4, 0):
+        inclusive_par = {'inclusive': inclusive}
+    else:
+        if inclusive == 'neither':
+            inclusive = 'right'
+            df[end_date_col] -= pd.DateOffset(microseconds=1)
+        inclusive_par = {'closed': None if inclusive == 'both' else inclusive}
+
+    # if inclusive = 'left' we expect date_end is exclusive
+    #   so records with start_date == end_date should be excluded
+    # also reset index to ensure index start from 0 and is consecutive
+    if inclusive == 'left':
+        df = df.query(f'{start_date_col} < {end_date_col}').reset_index(drop=True)
+    else:
+        df = df.query(f'{start_date_col} <= {end_date_col}').reset_index(drop=True)
+
+    # get exploded timestamp column
+    if df.empty:
+        dt = (
+            pd.DataFrame(
+                {
+                    'i': pd.Series(dtype='int'),
+                    date_col: pd.Series(dtype='datetime64[ns]'),
+                }
+            )
+            .set_index('i')
+            .rename_axis(None, axis=0)
+        )
+    else:
+        dt = (
+            pd.concat(
+                [
+                    pd.DataFrame(
+                        {
+                            'i': i,
+                            date_col: pd.date_range(
+                                start=s, end=e, freq=freq, **inclusive_par
+                            ),
+                        }
+                    )
+                    for i, (s, e) in enumerate(
+                        zip(df[start_date_col], df[end_date_col])
+                    )
+                ]
+            )
+            .set_index('i')
+            .rename_axis(None, axis=0)
+        )
+
+    # drop start_date_col and end_date_col
+    if drop_date_cols:
+        df = df.drop(columns=[start_date_col, end_date_col])
+
+    # sample df based on new timestamp column
+    df = df.reindex(dt.index)
+    df[date_col] = dt[date_col]
+
+    # set index
+    if drop_index:
+        df = df.reset_index(drop=True)
+    else:
+        df = df.set_index(levels).rename_axis(levels_old, axis=0)
+
+    return df

mspu/pandas/parquet.py

@@ -0,0 +1,36 @@
+import pandas as pd
+
+
+def pa_mod(
+    ds: int | pd.Series,
+    divisor: int,
+) -> int | pd.Series:
+    """
+    Calculates remainder after division by a positive divisor.
+
+    Parameters:
+    -----------
+    ds: int | pd.Series
+        An integer value or a pd.Series with int64[pyarrow] dtype
+    divisor: int
+        The positive divisor for the modulo operation
+
+    Returns:
+    -----------
+    int | pd.Series:
+        An integer or pd.series containing the modulo results
+    """
+
+    if divisor <= 0:
+        raise ValueError('Divisor must be a positive integer')
+
+    # Check if divisor is a power of 2
+    if divisor & (divisor - 1) == 0:
+        # Efficient bitwise AND for power-of-2 divisors
+        remainder = ds & (divisor - 1)
+    else:
+        # Slower integer division for non-power-of-2 divisors
+        quotient = ds // divisor
+        remainder = ds - (quotient * divisor)
+
+    return remainder

mspu/pandas/utils.py

@@ -0,0 +1,196 @@
+import pandas as pd
+
+
+def pd_ht(self, n: int = 2, c: int = None, w: int = None, r: int = None) -> None:
+    """
+    Pandas head and tail in one command, with optional rounding.
+
+    Parameters
+    ----------
+    n : int
+        Number of rows to show from the head and tail. If n < 0 or n is greater than
+        half the number of rows, the entire DataFrame will be shown.
+    c : int
+        Number of columns to show. If None or greater than the number of columns,
+        all columns will be shown.
+    w : int
+        Display width in characters. If None, pandas default is used.
+    r : int
+        Number of decimal places to round float columns. If None or negative,
+        no rounding will be applied.
+
+    Returns
+    -------
+    None
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> from mspu.pandas.utils import pd_ht
+    >>> pd.DataFrame.ht = pd_ht
+    >>> df = pd.DataFrame({
+    ...     'foo': [1.12345, 2.98765, 3.14159],
+    ...     'bar': [7, 8, 9],
+    ...     'ham': ['x', 'y', 'z'],
+    ... })
+    >>> df.ht(n=1, c=2, r=2)
+    shape: (3, 3)
+        foo  ...  ham
+    0  1.12  ...    x
+    2  3.14  ...    z
+    """
+    # Columns to show
+    if isinstance(c, int) and c <= 0:
+        c = None
+
+    # Display width
+    if isinstance(w, int) and w <= 0:
+        w = 1000
+
+    # Slice head + tail (or full df)
+    if n < 0 or self.shape[0] <= 2 * n:
+        df = self
+    else:
+        df = pd.concat([self.iloc[:n], self.iloc[-n:]])
+
+    # Round float columns
+    if r is not None and r >= 0:
+        float_cols = df.select_dtypes(include='float').columns
+        df[float_cols] = df[float_cols].round(r)
+
+    with pd.option_context(
+        'display.show_dimensions',
+        False,
+        'display.width',
+        w,
+        'display.max_rows',
+        df.shape[0],
+        'display.max_columns',
+        c,
+    ):
+        print(f'shape: {self.shape}')
+        print(df)
+
+
+def remove_cols_utc(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Converts datetime columns in UTC to timezone-naive (tz=None).
+    """
+    for col in df.select_dtypes(include=['datetimetz']).columns:
+        tz = df[col].dt.tz
+        if tz is not None and 'utc' in str(tz).lower():
+            df[col] = df[col].dt.tz_localize(None)
+    return df
+
+
+def df_diffs(
+    df1: pd.DataFrame,
+    df2: pd.DataFrame,
+    left_suffix: str = '_df1',
+    right_suffix: str = '_df2',
+) -> pd.DataFrame:
+    """
+    Get rows in two dfs that are different to each other.
+
+    Comparison is based on df MultiIndex.
+
+    Parameters
+    ----------
+    df1 : pd.DataFrame
+        First DataFrame to compare.
+    df2 : pd.DataFrame
+        Second DataFrame to compare.
+    left_suffix : str
+        Suffix to append to column names from df1 in the output DataFrame.
+    right_suffix : str
+        Suffix to append to column names from df2 in the output DataFrame.
+
+    Returns
+    -------
+    pd.DataFrame
+        DataFrame containing rows that are different between df1 and df2, with suffixes
+        indicating the source of each column.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> from mspu.pandas import df_diffs
+    >>> df1 = pd.DataFrame({
+    ...     'fruit': ['apple', 'banana', 'apple'],
+    ...     'id': [1, 2, 4],
+    ...     'store': ['us', 'uk', 'cn'],
+    ...     'price': [4, 3.1, 2.5],
+    ... }).set_index(['fruit', 'id'])
+    >>> print(df1)
+            store  price
+    fruit  id
+    apple  1     us    4.0
+    banana 2     uk    3.1
+    apple  4     cn    2.5
+    >>> df2 = pd.DataFrame({
+    ...     'fruit': ['apple', 'banana', 'apple'],
+    ...     'id': [1, 3, 4],
+    ...     'store': ['us', 'uk', 'cn'],
+    ...     'price': [4, 3.2, 2.5],
+    ... }).set_index(['fruit', 'id'])
+    >>> print(df2)
+            store  price
+    fruit  id
+    apple  1     us    4.0
+    banana 3     uk    3.2
+    apple  4     cn    2.5
+    >>> diff = df_diffs(df1, df2)
+    >>> print(diff)
+              store_df1  price_df1 store_df2  price_df2
+    fruit  id
+    banana 2         uk        3.1       NaN        NaN
+           3        NaN        NaN        uk        3.2
+    """
+    # Perform an outer join
+    df_joined = pd.merge(
+        df1,
+        df2,
+        how='outer',
+        left_index=True,
+        right_index=True,
+        suffixes=(left_suffix, right_suffix),
+    )
+
+    # Compare the columns to find differences
+    d1_joined = df_joined.filter(like=left_suffix).rename(
+        columns=lambda x: x.rsplit('_', 1)[0]
+    )
+    d2_joined = df_joined.filter(like=right_suffix).rename(
+        columns=lambda x: x.rsplit('_', 1)[0]
+    )
+
+    # Compare two dfs with `ne`, col names must be the same
+    diffs = df_joined[d1_joined.ne(d2_joined).any(axis=1)]
+
+    return diffs
+
+
+if __name__ == '__main__':
+    # Create two example DataFrames with string columns
+    df1 = pd.DataFrame(
+        {
+            'fruit': ['apple', 'banana', 'apple'],
+            'id': [1, 2, 4],
+            'store': ['us', 'uk', 'cn'],
+            'price': [4, 3.1, 2.5],
+        }
+    ).set_index(['fruit', 'id'])
+
+    df2 = pd.DataFrame(
+        {
+            'fruit': ['apple', 'banana', 'apple'],
+            'id': [1, 3, 4],
+            'store': ['us', 'uk', 'cn'],
+            'price': [4, 3.2, 2.5],
+        }
+    ).set_index(['fruit', 'id'])
+
+    df = df_diffs(df1, df2)
+
+    # Show the resulting differences
+    print(df)

mspu/polars/__init__.py

@@ -0,0 +1,17 @@
+from .utils import (
+    pl_ht,
+    inf_count,
+    nan_count,
+    nul_count,
+    lowercase_polars_df,
+    to_float32_polars_df,
+)
+
+__all__ = [
+    'pl_ht',
+    'inf_count',
+    'nan_count',
+    'nul_count',
+    'lowercase_polars_df',
+    'to_float32_polars_df',
+]

mspu/polars/utils.py

@@ -0,0 +1,142 @@
+import polars as pl
+import polars.selectors as cs
+from typing import Literal
+
+
+def pl_ht(self, n: int = 2, c: int = None, w: int = None, r: int = None) -> None:
+    """
+    Polars head and tail in one command, with optional rounding.
+
+    Parameters
+    ----------
+    n : int
+        Number of rows to show from the head and tail. If n < 0 or n is greater than
+        half the number of rows, the entire DataFrame will be shown.
+    c : int
+        Number of columns to show. If None or greater than the number of columns,
+        all columns will be shown.
+    w : int
+        Width of the output in characters. If None, the width will be determined.
+    r : int
+        Number of decimal places to round float and decimal columns. If None or negative,
+        no rounding will be applied.
+
+    Returns
+    -------
+    None
+
+    Examples
+    --------
+    >>> import polars as pl
+    >>> from mspu.polars.utils import pl_ht
+    >>> pl.DataFrame.ht = pl_ht
+    >>> df = pl.DataFrame({
+    ...   'foo': [1.12345, 2.98765, 3.14159],
+    ...   'bar': [7, 8, 9],
+    ...   'ham': ['x', 'y', 'z'],
+    ... })
+    >>> df.ht(n=1, c=2, r=2)
+    shape: (3, 3)
+    ┌──────┬───┬─────┐
+    │ foo  ┆ … ┆ ham │
+    │ ---  ┆   ┆ --- │
+    │ f64  ┆   ┆ str │
+    ╞══════╪═══╪═════╡
+    │ 1.12 ┆ … ┆ x   │
+    │ 3.14 ┆ … ┆ z   │
+    └──────┴───┴─────┘
+    """
+    if n < 0 or self.shape[0] < 2 * n:
+        df = self
+    else:
+        df = pl.concat([self[:n], self[-n:]])
+    if r is not None and r >= 0:
+        df = df.with_columns((cs.float() | cs.decimal()).round(r))
+    with pl.Config(
+        tbl_hide_dataframe_shape=True,
+        tbl_width_chars=w,
+        tbl_rows=df.shape[0],
+        tbl_cols=c,
+    ):
+        print(f'shape: {self.shape}')
+        print(df)
+
+
+def parquet_to_csv(filepath: str):
+    """
+    Read parquet and save as csv.
+    """
+    filepath_csv = f'{filepath[:-7]}csv'
+    pl.read_parquet(filepath).write_csv(filepath_csv)
+
+
+def lowercase_polars_df(
+    df: pl.DataFrame,
+    lowercase: Literal['header', 'columns', 'both'] = 'both',
+) -> pl.DataFrame:
+    """
+    Converts all column names and string columns to lowercase.
+    """
+    # Lowercase column headers
+    if lowercase in ('header', 'both'):
+        df = df.rename({col: col.lower() for col in df.columns})
+    # Lowercase string columns
+    if lowercase in ('columns', 'both'):
+        df = df.with_columns([cs.string().str.to_lowercase()])
+    return df
+
+
+def to_float32_polars_df(df: pl.DataFrame) -> pl.DataFrame:
+    """
+    Convert all numerical columns type to float32.
+    """
+    df = df.with_columns((cs.float() | cs.decimal()).cast(pl.Float32))
+    return df
+
+
+def inf_count(df: pl.DataFrame) -> pl.DataFrame:
+    """
+    Counts the number of infinite values in each column of a Polars DataFrame.
+
+    Returns a new DataFrame with the column names and their corresponding counts,
+    sorted in descending order.
+    """
+    df_inf = (
+        df.select([(pl.col(col).is_infinite().sum().alias(col)) for col in df.columns])
+        .unpivot(variable_name='col', value_name='inf_cnt')
+        .filter(pl.col('inf_cnt') > 0)
+        .sort('inf_cnt', descending=True)
+    )
+    return df_inf
+
+
+def nan_count(df: pl.DataFrame) -> pl.DataFrame:
+    """
+    Counts the number of NaN values in each column of a Polars DataFrame.
+
+    Returns a new DataFrame with the column names and their corresponding counts,
+    sorted in descending order.
+    """
+    df_nan = (
+        df.select([(pl.col(col).is_nan().sum().alias(col)) for col in df.columns])
+        .unpivot(variable_name='col', value_name='nan_cnt')
+        .filter(pl.col('nan_cnt') > 0)
+        .sort('nan_cnt', descending=True)
+    )
+    return df_nan
+
+
+def nul_count(df: pl.DataFrame) -> pl.DataFrame:
+    """
+    Counts the number of null values in each column of a Polars DataFrame.
+
+    Returns a new DataFrame with the column names and their corresponding counts,
+    sorted in descending order.
+    """
+    df_nul = (
+        df.select([(pl.col(col).is_null().sum().alias(col)) for col in df.columns])
+        .unpivot(variable_name='col', value_name='nul_cnt')
+        .filter(pl.col('nul_cnt') > 0)
+        .sort('nul_cnt', descending=True)
+    )
+    return df_nul

mspu/security/crypt.py

@@ -0,0 +1,80 @@
+import json
+import base64
+import hashlib
+from cryptography.fernet import Fernet
+
+
+def create_key(keypath: str):
+    # Generate a key
+    key = Fernet.generate_key()
+
+    # Save the key securely
+    with open(keypath, 'wb') as f:
+        f.write(key)
+
+
+def load_key(keypath: str) -> bytes:
+    # Load key from file
+    with open(keypath, 'rb') as f:
+        key = f.read()
+    return key
+
+
+def encrypt_txt(key: bytes, txt: str) -> bytes:
+    # Serialize and encrypt the txt
+    cipher = Fernet(key)
+    encrypted_credentials = cipher.encrypt(txt)
+    return encrypted_credentials
+
+
+def decrypt_txt(key: bytes, txt: bytes) -> bytes:
+    # Decrypt the credentials
+    cipher = Fernet(key)
+    decrypted_credentials = cipher.decrypt(txt)
+    return decrypted_credentials
+
+
+def encrypt_file(key: bytes, filepath: str):
+    # Load the credentials
+    with open(filepath, 'rb') as f:
+        txt = f.read()
+    # Serialize and encrypt the txt
+    encrypted_credentials = encrypt_txt(key, txt)
+    # Save the encrypted file
+    file, ext = filepath.rsplit('.', 1)
+    filepath = f'{file}_.{ext}'
+    with open(filepath, 'wb') as f:
+        f.write(encrypted_credentials)
+
+
+def decrypt_file(key: bytes, filepath: str, to_json: bool = True):
+    # Load the credentials
+    with open(filepath, 'rb') as f:
+        txt = f.read()
+    # Decrypt the credentials
+    decrypted_credentials = decrypt_txt(key, txt).decode()
+    # Load the credentials into a dictionary
+    if to_json:
+        decrypted_credentials = json.loads(decrypted_credentials)
+    return decrypted_credentials
+
+
+def ha64(txt: str):
+    txt = hashlib.sha256(txt.encode()).digest()
+    return base64.urlsafe_b64encode(txt)
+
+if __name__ == '__main__':
+    fp = 'c:/files'
+    # create_key(f'{fp}/my-key.json')
+
+    key = load_key(f'{fp}/my-key.json')
+
+    key = hashlib.sha256('my-key'.encode()).digest()
+    key = base64.urlsafe_b64encode(key)
+    fl = 'my-key.json'
+    # encrypt_file(key, f'{fp}/{fl}')
+
+    key = decrypt_file(key, f'{fp}/my-key_.json', False)
+    db_cred = decrypt_file(key, f'{fp}/credentials_.json')
+
+    print(db_cred)

mspu-0.0.1.dist-info/METADATA

@@ -0,0 +1,32 @@
+Metadata-Version: 2.4
+Name: mspu
+Version: 0.0.1
+Summary: Ma's python utils package.
+Author-email: Sean Ma <sean.sl.ma@gmail.com>
+License: MIT
+License-File: LICENSE
+Requires-Python: >=3.12
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: polars
+Description-Content-Type: text/markdown
+
+# mspu
+
+My python utils package. It includes some commonly used functions, like `gen_rand_df` and `explode_date_range`.
+
+## How to install `mspu`
+
+Using `pip`:
+```sh
+pip install mspu
+```
+
+Using `uv`:
+```sh
+uv pip install mspu
+```
+
+## The document for `mspu`
+
+Visit the [document](https://seanslma.github.io/mspu) for more details.

mspu-0.0.1.dist-info/RECORD

@@ -0,0 +1,18 @@
+mspu/__init__.py,sha256=uWOTnD7uFq1Qca0T90qOTjvIq5on-r64Ug4hJWPYQ64,184
+mspu/data/__init__.py,sha256=THICntQKY0THohm-I1gPwi9I0yJHHg5Jv0IVCgZ8uAA,67
+mspu/data/rand_df.py,sha256=I1Hkc8NfrLamff2fUG6il0KW67epoVrvq9TrKJSVS6w,10841
+mspu/datetime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mspu/datetime/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mspu/io/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mspu/io/file.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mspu/pandas/__init__.py,sha256=gommfn2jXrmUBnXG2qMtQzqRMg1oEqyLPMGmRzGSgmA,188
+mspu/pandas/datetime.py,sha256=_iAZ9xqFJ-tYnT3x5hT09wt0Ihlc21PPaBOfFyL8boc,8656
+mspu/pandas/parquet.py,sha256=XjjdpgNG7Qsm3KbpbuR1qfGKFnq1FdPaoJzmA19A7I8,915
+mspu/pandas/utils.py,sha256=3JzKQDssYzxJosgtARc90EFQCb24yHqyNmGJlWx6dDk,5542
+mspu/polars/__init__.py,sha256=dh4ekg-o_Gv4-nnnYI3flRliNn9BF16fJJE97_vCBzI,264
+mspu/polars/utils.py,sha256=VMZrljTV7mW7Uy0SpMRS2aV6whhTIgDMYrRNbHf7bIA,4511
+mspu/security/crypt.py,sha256=iYnkg9YdPJjEgYrw7iAnpK0qiH0illpp1-bIkZpIxF0,2097
+mspu-0.0.1.dist-info/METADATA,sha256=ZHGLKpsG-e2WwuBTpCxzVpOT9hmBDHKFgRqJQC9BKd8,625
+mspu-0.0.1.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+mspu-0.0.1.dist-info/licenses/LICENSE,sha256=Rc482qeBCHvWChKavd4nCgfkKshLEHDcmTlXNOIQ6DY,1064
+mspu-0.0.1.dist-info/RECORD,,

mspu-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any

mspu-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Sean Ma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.