oafuncs 0.0.59__py2.py3-none-any.whl

oafuncs/oa_file.py

@@ -0,0 +1,297 @@
+#!/usr/bin/env python
+# coding=utf-8
+'''
+Author: Liu Kun && 16031215@qq.com
+Date: 2024-09-17 15:07:13
+LastEditors: Liu Kun && 16031215@qq.com
+LastEditTime: 2024-11-11 16:11:23
+FilePath: \\Python\\My_Funcs\\OAFuncs\\OAFuncs\\oa_file.py
+Description:
+EditPlatform: vscode
+ComputerInfo: XPS 15 9510
+SystemInfo: Windows 11
+Python Version: 3.11
+'''
+
+
+import glob
+import os
+import re
+import shutil
+
+__all__ = ['link_file', 'copy_file', 'rename_files', 'make_folder',
+           'clear_folder', 'remove_empty_folders', 'remove', 'file_size']
+
+
+def link_file(src_pattern, dst):
+    '''
+    # 描述：创建符号链接，支持通配符
+    # 使用示例
+    # link_file(r'/data/hejx/liukun/era5/*', r'/data/hejx/liukun/Test/')
+    # link_file(r'/data/hejx/liukun/era5/py.o*', r'/data/hejx/liukun/Test/py.o')
+    # link_file(r'/data/hejx/liukun/era5/py.o*', r'/data/hejx/liukun/Test')
+    param      {*} src_pattern # 源文件或目录
+    param      {*} dst # 目标文件或目录
+    '''
+    src_pattern= str(src_pattern)
+    # 使用glob.glob来处理可能包含通配符的src
+    src_files= glob.glob(src_pattern)
+    if not src_files:
+        raise FileNotFoundError('文件不存在: {}'.format(src_pattern))
+
+    # 判断dst是路径还是包含文件名的路径
+    if os.path.isdir(dst):
+        # 如果dst是路径，则保持源文件的文件名
+        dst_dir= dst
+        for src_file in src_files:
+            src_file_basename= os.path.basename(src_file)
+            dst_file= os.path.join(dst_dir, src_file_basename)
+            if os.path.exists(dst_file):
+                os.remove(dst_file)
+            os.symlink(src_file, dst_file)
+            print(f'创建符号链接: {src_file} -> {dst_file}')
+    else:
+        # 如果dst包含文件名，则创建链接后重命名
+        dst_dir= os.path.dirname(dst)
+        os.makedirs(dst_dir, exist_ok=True)
+        # 只处理第一个匹配的文件
+        src_file= src_files[0]
+        dst_file= dst
+        if os.path.exists(dst_file):
+            os.remove(dst_file)
+        os.symlink(src_file, dst_file)
+        print(f'创建符号链接并重命名: {src_file} -> {dst_file}')
+
+
+def copy_file(src_pattern, dst):
+    '''
+    # 描述：复制文件或目录，支持通配符
+    # 使用示例
+    # copy_file(r'/data/hejx/liukun/era5/py.o*', r'/data/hejx/liukun/Test/py.o')
+    # copy_file(r'/data/hejx/liukun/era5/py.o*', r'/data/hejx/liukun/Test')
+    param      {*} src_pattern # 源文件或目录
+    param      {*} dst # 目标文件或目录
+    '''
+    src_pattern= str(src_pattern)
+    # 使用glob.glob来处理可能包含通配符的src
+    src_files= glob.glob(src_pattern)
+    if not src_files:
+        raise FileNotFoundError('文件不存在: {}'.format(src_pattern))
+
+    # 判断dst是路径还是包含文件名的路径
+    if os.path.isdir(dst):
+        # 如果dst是路径，则保持源文件的文件名
+        dst_dir= dst
+        for src_file in src_files:
+            src_file_basename= os.path.basename(src_file)
+            dst_file= os.path.join(dst_dir, src_file_basename)
+            if os.path.exists(dst_file):
+                if os.path.isdir(dst_file):
+                    shutil.rmtree(dst_file)
+                else:
+                    os.remove(dst_file)
+            if os.path.isdir(src_file):
+                shutil.copytree(src_file, dst_file, symlinks=True)
+            else:
+                shutil.copy2(src_file, dst_file)
+            print(f'复制文件或目录: {src_file} -> {dst_file}')
+    else:
+        # 如果dst包含文件名，则复制后重命名
+        dst_dir= os.path.dirname(dst)
+        os.makedirs(dst_dir, exist_ok=True)
+        # 只处理第一个匹配的文件
+        src_file= src_files[0]
+        dst_file= dst
+        if os.path.exists(dst_file):
+            if os.path.isdir(dst_file):
+                shutil.rmtree(dst_file)
+            else:
+                os.remove(dst_file)
+        if os.path.isdir(src_file):
+            shutil.copytree(src_file, dst_file, symlinks=True)
+        else:
+            shutil.copy2(src_file, dst_file)
+        print(f'复制文件或目录并重命名: {src_file} -> {dst_file}')
+
+
+def rename_files(directory, old_str, new_str):
+    '''
+    # 描述：重命名目录下的文件，支持通配符
+    # 使用示例
+    directory_path = r"E:\windfarm\CROCO_FILES"
+    old_str = "croco"
+    new_str = "roms"
+    rename_files(directory_path, old_str, new_str)
+    param      {*} directory # 目录
+    param      {*} old_str # 要替换的字符串
+    param      {*} new_str # 新字符串
+    '''
+    # 获取目录下的所有文件
+    files= os.listdir(directory)
+
+    # 构建正则表达式以匹配要替换的字符串
+    pattern= re.compile(re.escape(old_str))
+
+    # 遍历目录下的文件
+    for filename in files:
+        # 检查文件名中是否包含要替换的字符串
+        if pattern.search(filename):
+            # 构建新的文件名
+            new_filename= pattern.sub(new_str, filename)
+
+            # 构建旧文件的完整路径
+            old_path= os.path.join(directory, filename)
+
+            # 构建新文件的完整路径
+            new_path= os.path.join(directory, new_filename)
+
+            # 重命名文件
+            os.rename(old_path, new_path)
+            print(f"重命名文件: {old_path} -> {new_path}")
+
+
+# ** 创建子文件夹（可选清空）
+def make_folder(rootpath: str, folder_name: str, clear=0) -> str:
+    '''
+    # 描述：创建子文件夹（可选清空）
+    # 使用示例
+    rootpath = r'E:\Data\2024\09\17'
+    folder_name = 'var1'
+    newpath = make_folder(rootpath, folder_name, clear=1)
+    param       {*} rootpath # 根目录
+    param       {*} folder_name # 文件夹名称
+    '''
+    folder_path= os.path.join(str(rootpath), str(folder_name))
+    if clear:
+        shutil.rmtree(folder_path, ignore_errors=True)
+    os.makedirs(folder_path, exist_ok=True)
+    return folder_path
+
+# ** 清空文件夹
+
+
+def clear_folder(folder_path):
+    '''
+    # 描述：清空文件夹
+    # 使用示例
+    clear_folder(r'E:\Data\2024\09\17\var1')
+    param        {*} folder_path # 文件夹路径
+    '''
+    folder_path= str(folder_path)
+    if os.path.exists(folder_path):
+        try:
+            # 遍历文件夹中的所有文件和子文件夹
+            for filename in os.listdir(folder_path):
+                file_path= os.path.join(folder_path, filename)
+                # 判断是文件还是文件夹
+                if os.path.isfile(file_path) or os.path.islink(file_path):
+                    os.unlink(file_path)  # 删除文件或链接
+                elif os.path.isdir(file_path):
+                    shutil.rmtree(file_path)  # 删除子文件夹
+            print(f'成功清空文件夹: {folder_path}')
+        except Exception as e:
+            print(f'清空文件夹失败: {folder_path}')
+            print(e)
+
+
+# ** 清理空文件夹
+def remove_empty_folders(path, print_info=1):
+    '''
+    # 描述：清理空文件夹
+    # 使用示例
+    remove_empty_folders(r'E:\Data\2024\09\17', print_info=1)
+    param        {*} path # 文件夹路径
+    param        {*} print_info # 是否打印信息
+    '''
+    path= str(path)
+    # 遍历当前目录下的所有文件夹和文件
+    for root, dirs, files in os.walk(path, topdown=False):
+        # 遍历文件夹列表
+        for folder in dirs:
+            folder_path= os.path.join(root, folder)
+            # 判断文件是否有权限访问
+            try:
+                os.listdir(folder_path)
+            except OSError:
+                continue
+            # 判断文件夹是否为空
+            if not os.listdir(folder_path):
+                # 删除空文件夹
+                try:
+                    os.rmdir(folder_path)
+                    print(f"Deleted empty folder: {folder_path}")
+                except OSError:
+                    if print_info:
+                        print(f"Skipping protected folder: {folder_path}")
+                    pass
+
+
+# ** 删除相关文件，可使用通配符
+def remove(pattern):
+    '''
+    # 描述：删除相关文件，可使用通配符
+    remove(r'E:\Code\Python\Model\WRF\Radar2\bzip2-radar-0*')
+    # or
+    os.chdir(r'E:\Code\Python\Model\WRF\Radar2')
+    remove('bzip2-radar-0*')
+    param        {*} pattern # 文件路径或通配符
+    '''
+    # 使用glob.glob来获取所有匹配的文件
+    # 可以使用通配符*来匹配所有文件
+    pattern= str(pattern)
+    file_list= glob.glob(pattern)
+    for file_path in file_list:
+        if os.path.exists(file_path):
+            try:
+                shutil.rmtree(file_path)
+                print(f'成功删除文件: {file_path}')
+            except Exception as e:
+                print(f'删除文件失败: {file_path}')
+                print(e)
+        else:
+            print(f'文件不存在: {file_path}')
+
+
+def file_size(file_path, unit='KB'):
+    '''
+    Description: 获取文件大小
+
+    Args:
+    file_path: 文件路径
+    unit: 单位（PB、TB、GB、MB、KB）
+
+    Returns:
+    文件大小（单位：PB、TB、GB、MB、KB）
+    '''
+    # 检查文件是否存在
+    if not os.path.exists(file_path):
+        return "文件不存在"
+
+    # 获取文件大小（字节）
+    file_size= os.path.getsize(file_path)
+
+    # 单位转换字典
+    unit_dict= {
+        'PB': 1024**5,
+        'TB': 1024**4,
+        'GB': 1024**3,
+        'MB': 1024**2,
+        'KB': 1024
+    }
+
+    # 检查传入的单位是否合法
+    if unit not in unit_dict:
+        return "单位不合法，请选择PB、TB、GB、MB、KB中的一个"
+
+    # 转换文件大小到指定单位
+    converted_size= file_size / unit_dict[unit]
+
+    return converted_size
+
+
+if __name__ == '__main__':
+    # newpath = make_folder('D:/Data/2024/09/17/', 'var1', clear=1)
+    # print(newpath)
+    pass
+
+    remove(r'I:\Delete\test\*')

oafuncs/oa_help.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+# coding=utf-8
+'''
+Author: Liu Kun && 16031215@qq.com
+Date: 2024-10-06 19:25:29
+LastEditors: Liu Kun && 16031215@qq.com
+LastEditTime: 2024-10-06 19:55:04
+FilePath: \\Python\\My_Funcs\\OAFuncs\\OAFuncs\\oa_help.py
+Description:  
+EditPlatform: vscode
+ComputerInfo: XPS 15 9510
+SystemInfo: Windows 11
+Python Version: 3.11
+'''
+
+import OAFuncs
+
+
+def query():
+    '''
+    description: 查看OAFuncs模块的函数列表
+    example: query()
+    '''
+    funcs = [func for func in dir(OAFuncs) if callable(getattr(OAFuncs, func))]
+    print('函数数量：')
+    print(len(funcs))
+    print('函数列表：')
+    print(funcs)
+
+
+def use(func='get_var'):
+    '''
+    description: 查看函数的模块全路径和函数提示
+    param {func} : 函数名
+    example: use('get_var')
+    '''
+    print('模块全路径：')
+    print(getattr(OAFuncs, func).__module__+'.'+func)
+    print('函数提示：')
+    print(getattr(OAFuncs, func).__doc__)

oafuncs/oa_nc.py

@@ -0,0 +1,345 @@
+#!/usr/bin/env python
+# coding=utf-8
+'''
+Author: Liu Kun && 16031215@qq.com
+Date: 2024-09-17 14:58:50
+LastEditors: Liu Kun && 16031215@qq.com
+LastEditTime: 2024-11-05 16:44:20
+FilePath: \\Python\\My_Funcs\\OAFuncs\\OAFuncs\\oa_nc.py
+Description:  
+EditPlatform: vscode
+ComputerInfo: XPS 15 9510
+SystemInfo: Windows 11
+Python Version: 3.11
+'''
+
+
+import os
+
+import netCDF4 as nc
+import numpy as np
+import xarray as xr
+
+__all__ = ['get_var', 'extract5nc', 'write2nc',
+           'merge5nc', 'modify_var_value', 'modify_var_attr', 'rename_var_or_dim', 'check_ncfile']
+
+
+def get_var(file, *vars):
+    '''
+    description: 读取nc文件中的变量
+    param {file: 文件路径, *vars: 变量名}
+    example: datas = get_var(file_ecm, 'h', 't', 'u', 'v')
+    return {datas: 变量数据}
+    '''
+    ds = xr.open_dataset(file)
+    datas = []
+    for var in vars:
+        data = ds[var]
+        datas.append(data)
+    ds.close()
+    return datas
+
+
+def extract5nc(file, varname):
+    '''
+    描述：
+    1、提取nc文件中的变量
+    2、将相应维度提取，建立字典
+    return：返回变量及坐标字典
+    参数：
+    file: 文件路径
+    varname: 变量名
+    example: data, dimdict = extract5nc(file_ecm, 'h')
+    '''
+    ds = xr.open_dataset(file)
+    vardata = ds[varname]
+    dims = vardata.dims
+    dimdict = {}
+    for dim in dims:
+        dimdict[dim] = ds[dim].values
+    ds.close()
+    return np.array(vardata), dimdict
+
+
+def _numpy_to_nc_type(numpy_type):
+    """将NumPy数据类型映射到NetCDF数据类型"""
+    numpy_to_nc = {
+        'float32': 'f4',
+        'float64': 'f8',
+        'int8': 'i1',
+        'int16': 'i2',
+        'int32': 'i4',
+        'int64': 'i8',
+        'uint8': 'u1',
+        'uint16': 'u2',
+        'uint32': 'u4',
+        'uint64': 'u8',
+    }
+    return numpy_to_nc.get(str(numpy_type), 'f4')  # 默认使用 'float32'
+
+
+def write2nc(file, data, varname, coords, mode):
+    '''
+    description: 写入数据到nc文件
+    参数：
+    file: 文件路径
+    data: 数据
+    varname: 变量名
+    coords: 坐标，字典，键为维度名称，值为坐标数据
+    mode: 写入模式，'w'为写入，'a'为追加
+    example: write2nc(r'test.nc', data, 'data', {'time': np.linspace(0, 120, 100), 'lev': np.linspace(0, 120, 50)}, 'a')
+    '''
+    # 判断mode是写入还是追加
+    if mode == 'w':
+        if os.path.exists(file):
+            os.remove(file)
+            print("Warning: File already exists. Deleting it.")
+    elif mode == 'a':
+        if not os.path.exists(file):
+            print("Warning: File doesn't exist. Creating a new file.")
+            mode = 'w'
+
+    # 打开 NetCDF 文件
+    with nc.Dataset(file, mode, format="NETCDF4") as ncfile:
+        # 处理坐标
+        for dim, coord_data in coords.items():
+            add_coords = True
+            # 判断坐标是否存在，若存在，则替换/报错
+            if ncfile.dimensions:
+                # 返回字典，字典、列表、元组若为空，都表示False
+                if dim in ncfile.dimensions:
+                    # del nc.dimensions[dim]
+                    if len(coord_data) != len(ncfile.dimensions[dim]):
+                        raise ValueError(
+                            "Length of coordinate does not match the dimension length.")
+                    else:
+                        add_coords = False
+                        print(
+                            f"Warning: Coordinate '{dim}' already exists. Replacing it.")
+                        ncfile.variables[dim][:] = np.array(coord_data)
+            if add_coords:
+                # 创建新坐标
+                ncfile.createDimension(dim, len(coord_data))
+                ncfile.createVariable(dim, _numpy_to_nc_type(
+                    coord_data.dtype), (dim,))
+                ncfile.variables[dim][:] = np.array(coord_data)
+
+        # 判断变量是否存在，若存在，则删除原变量
+        add_var = True
+        if varname in ncfile.variables:
+            print(f"Warning: Variable '{varname}' already exists.")
+            if data.shape != ncfile.variables[varname].shape:
+                raise ValueError(
+                    "Shape of data does not match the variable shape.")
+            else:
+                # 写入数据
+                ncfile.variables[varname][:] = data
+                add_var = False
+                print(
+                    f"Warning: Variable '{varname}' already exists. Replacing it.")
+
+        if add_var:
+            # 创建变量及其维度
+            dim_names = tuple(coords.keys())  # 使用coords传入的维度名称
+            ncfile.createVariable(
+                varname, _numpy_to_nc_type(data.dtype), dim_names)
+            # ncfile.createVariable('data', 'f4', ('time','lev'))
+
+            # 写入数据
+            ncfile.variables[varname][:] = data
+
+        # 判断维度是否匹配
+        if len(data.shape) != len(coords):
+            raise ValueError(
+                "Number of dimensions does not match the data shape.")
+
+
+def merge5nc(file_list, var_name, dim_name, target_filename):
+    """
+    批量提取 nc 文件中的某一变量，按照某一维度合并后写入新的 nc 文件。
+
+    参数：
+    file_list：nc 文件路径列表
+    var_name：要提取的变量名
+    dim_name：用于合并的维度名
+    target_filename：合并后的目标文件名
+    example: merge5nc(file_list, 'sst', 'time', 'merged_data.nc')
+    """
+    data_list = []
+    for i, file in enumerate(file_list):
+        print(f"\rReading file {i + 1}/{len(file_list)}...", end="")
+        ds = xr.open_dataset(file)
+        var = ds[var_name]
+        data_list.append(var)
+        ds.close()
+    print("\nMerging data...")
+    data = xr.concat(data_list, dim=dim_name)
+    print("Writing data to file...")
+    if os.path.exists(target_filename):
+        print("Warning: The target file already exists.")
+        print("Removing existing file...")
+        os.remove(target_filename)
+    data.to_netcdf(target_filename)
+
+
+def merge5nc_vars(file_list, var_names, dim_name, target_filename):
+    """
+    批量提取 nc 文件中的两个变量，按照某一维度合并后写入新的 nc 文件。
+
+    参数：
+    file_list：nc 文件路径列表
+    var_names：要提取的变量名列表，例如 ['u', 'v']
+    dim_name：用于合并的维度名
+    target_filename：合并后的目标文件名
+
+    example: merge5nc_vars(file_list, ['u', 'v'], 'time', 'merged_data.nc')
+    """
+    data_lists = [[] for _ in var_names]
+    for i, file in enumerate(file_list):
+        print(f"\rReading file {i + 1}/{len(file_list)}...", end="")
+        ds = xr.open_dataset(file)
+        for j, var_name in enumerate(var_names):
+            var = ds[var_name]
+            data_lists[j].append(var)
+        ds.close()
+    print("\nMerging data...")
+    merged_data = {}
+    for var_name, data_list in zip(var_names, data_lists):
+        merged_data[var_name] = xr.concat(data_list, dim=dim_name)
+    print("Writing data to file...")
+    ds_merged = xr.Dataset(merged_data)
+    if os.path.exists(target_filename):
+        print("Warning: The target file already exists.")
+        print("Removing existing file...")
+        os.remove(target_filename)
+    ds_merged.to_netcdf(target_filename)
+
+
+def modify_var_value(nc_file_path, variable_name, new_value):
+    """
+    使用 netCDF4 库修改 NetCDF 文件中特定变量的值
+
+    参数：
+    nc_file_path (str): NetCDF 文件路径
+    variable_name (str): 要修改的变量名
+    new_value (numpy.ndarray): 新的变量值
+
+    example: modify_var_value('test.nc', 'data', np.random.rand(100, 50))
+    """
+    try:
+        # Open the NetCDF file
+        dataset = nc.Dataset(nc_file_path, 'r+')
+        # Get the variable to be modified
+        variable = dataset.variables[variable_name]
+        # Modify the value of the variable
+        variable[:] = new_value
+        dataset.close()
+        print(
+            f"Successfully modified variable {variable_name} in {nc_file_path}.")
+    except Exception as e:
+        print(
+            f"An error occurred while modifying variable {variable_name} in {nc_file_path}: {e}")
+
+
+def modify_var_attr(nc_file_path, variable_name, attribute_name, attribute_value):
+    """
+    使用 netCDF4 库添加或修改 NetCDF 文件中特定变量的属性。
+
+    参数：
+    nc_file_path (str): NetCDF 文件路径
+    variable_name (str): 要操作的变量名
+    attribute_name (str): 属性名
+    attribute_value (任意类型): 属性值
+    example: modify_var_attr('test.nc', 'data', 'long_name', 'This is a test variable.')
+    """
+    try:
+        ds = nc.Dataset(nc_file_path, 'r+')
+        if variable_name not in ds.variables:
+            raise ValueError(
+                f"Variable '{variable_name}' not found in the NetCDF file.")
+
+        variable = ds.variables[variable_name]
+        if attribute_name in variable.ncattrs():
+            print(
+                f"Warning: Attribute '{attribute_name}' already exists. Replacing it.")
+            variable.setncattr(attribute_name, attribute_value)
+        else:
+            print(f"Adding attribute '{attribute_name}'...")
+            variable.setncattr(attribute_name, attribute_value)
+
+        ds.close()
+    except Exception as e:
+        raise RuntimeError(f"An error occurred: {e}")
+
+
+def rename_var_or_dim(ncfile_path, old_name, new_name):
+    """
+    Rename a variable and/or dimension in a NetCDF file.
+
+    Parameters:
+    ncfile_path (str): The path to the NetCDF file.
+    old_name (str): The name of the variable or dimension to be renamed.
+    new_name (str): The new name for the variable or dimension.
+
+    example: rename_var_or_dim('test.nc', 'time', 'ocean_time')
+    """
+    try:
+        with nc.Dataset(ncfile_path, 'r+') as dataset:
+            # If the old name is not found as a variable or dimension, print a message
+            if old_name not in dataset.variables and old_name not in dataset.dimensions:
+                print(
+                    f"Variable or dimension {old_name} not found in the file.")
+
+            # Attempt to rename the variable
+            if old_name in dataset.variables:
+                dataset.renameVariable(old_name, new_name)
+                print(
+                    f"Successfully renamed variable {old_name} to {new_name}.")
+
+            # Attempt to rename the dimension
+            if old_name in dataset.dimensions:
+                # Check if the new dimension name already exists
+                if new_name in dataset.dimensions:
+                    raise ValueError(
+                        f"Dimension name {new_name} already exists in the file.")
+                dataset.renameDimension(old_name, new_name)
+                print(
+                    f"Successfully renamed dimension {old_name} to {new_name}.")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+def check_ncfile(ncfile, if_delete=False):
+    if not os.path.exists(ncfile):
+        return False
+
+    try:
+        with nc.Dataset(ncfile, 'r') as f:
+            # 确保f被使用，这里我们检查文件中变量的数量
+            if len(f.variables) > 0:
+                return True
+            else:
+                # 如果没有变量，我们可以认为文件是损坏的
+                raise ValueError("File is empty or corrupted.")
+    except OSError as e:
+        # 捕获文件打开时可能发生的OSError
+        print(f"An error occurred while opening the file: {e}")
+        if if_delete:
+            os.remove(ncfile)
+            print(f"File {ncfile} has been deleted.")
+        return False
+    except Exception as e:
+        # 捕获其他可能的异常
+        print(f"An unexpected error occurred: {e}")
+        if if_delete:
+            os.remove(ncfile)
+            print(f"File {ncfile} has been deleted.")
+        return False
+
+
+if __name__ == '__main__':
+
+    data = np.random.rand(100, 50)
+    write2nc(r'test.nc', data,
+             'data', {'time': np.linspace(0, 120, 100), 'lev': np.linspace(0, 120, 50)}, 'a')