PandasTA 源码解析（十九）（2）

# 从config模块中导入sample_data变量
from .config import sample_data
# 从context模块中导入pandas_ta模块
from .context import pandas_ta
# 从unittest模块中导入skip和TestCase类
from unittest import skip, TestCase
# 从pandas模块中导入DataFrame类
from pandas import DataFrame
# 定义测试TrendExtension类，继承自TestCase类
class TestTrendExtension(TestCase):
    # 设置测试类的类方法setUpClass，用于设置测试数据
    @classmethod
    def setUpClass(cls):
        cls.data = sample_data
    # 设置测试类的类方法tearDownClass，用于清理测试数据
    @classmethod
    def tearDownClass(cls):
        del cls.data
    # 设置测试方法的setUp，无操作
    def setUp(self): pass
    # 设置测试方法的tearDown，无操作
    def tearDown(self): pass
    # 测试ADX扩展功能
    def test_adx_ext(self):
        # 调用ADX扩展方法，将结果附加到数据框中
        self.data.ta.adx(append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后三列的列名
        self.assertEqual(list(self.data.columns[-3:]), ["ADX_14", "DMP_14", "DMN_14"])
    # 测试AMAT扩展功能
    def test_amat_ext(self):
        # 调用AMAT扩展方法，将结果附加到数据框中
        self.data.ta.amat(append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后两列的列名
        self.assertEqual(list(self.data.columns[-2:]), ["AMATe_LR_8_21_2", "AMATe_SR_8_21_2"])
    # 测试Aroon扩展功能
    def test_aroon_ext(self):
        # 调用Aroon扩展方法，将结果附加到数据框中
        self.data.ta.aroon(append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后三列的列名
        self.assertEqual(list(self.data.columns[-3:]), ["AROOND_14", "AROONU_14", "AROONOSC_14"])
    # 测试Chop扩展功能
    def test_chop_ext(self):
        # 调用Chop扩展方法，将结果附加到数据框中，设置ln参数为False
        self.data.ta.chop(append=True, ln=False)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名
        self.assertEqual(self.data.columns[-1], "CHOP_14_1_100")
        # 再次调用Chop扩展方法，将结果附加到数据框中，设置ln参数为True
        self.data.ta.chop(append=True, ln=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名
        self.assertEqual(self.data.columns[-1], "CHOPln_14_1_100")
    # 测试CKSP扩展功能
    def test_cksp_ext(self):
        # 调用CKSP扩展方法，将结果附加到数据框中，设置tvmode参数为False
        self.data.ta.cksp(tvmode=False, append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后两列的列名
        self.assertEqual(list(self.data.columns[-2:]), ["CKSPl_10_3_20", "CKSPs_10_3_20"])
    # 测试CKSP_TV扩展功能
    def test_cksp_tv_ext(self):
        # 调用CKSP扩展方法，将结果附加到数据框中，设置tvmode参数为True
        self.data.ta.cksp(tvmode=True, append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后两列的列名
        self.assertEqual(list(self.data.columns[-2:]), ["CKSPl_10_1_9", "CKSPs_10_1_9"])
    # 测试Decay扩展功能
    def test_decay_ext(self):
        # 调用Decay扩展方法，将结果附加到数据框中
        self.data.ta.decay(append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名
        self.assertEqual(self.data.columns[-1], "LDECAY_5")
        # 再次调用Decay扩展方法，设置mode参数为"exp"，将结果附加到数据框中
        self.data.ta.decay(mode="exp", append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名
        self.assertEqual(self.data.columns[-1], "EXPDECAY_5")
    # 测试Decreasing扩展功能
    def test_decreasing_ext(self):
        # 调用Decreasing扩展方法，将结果附加到数据框中
        self.data.ta.decreasing(append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名
        self.assertEqual(self.data.columns[-1], "DEC_1")
        # 再次调用Decreasing扩展方法，设置length参数为3，strict参数为True，将结果附加到数据框中
        self.data.ta.decreasing(length=3, strict=True, append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名
        self.assertEqual(self.data.columns[-1], "SDEC_3")
    # 测试DPO扩展功能
    def test_dpo_ext(self):
        # 调用DPO扩展方法，将结果附加到数据框中
        self.data.ta.dpo(append=True)
        # 断言数据类型为DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名
        self.assertEqual(self.data.columns[-1], "DPO_20")
    # 测试增长指标计算函数的扩展功能
    def test_increasing_ext(self):
        # 调用增长指标计算函数，将结果追加到原始数据后面
        self.data.ta.increasing(append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名为 "INC_1"
        self.assertEqual(self.data.columns[-1], "INC_1")
        # 调用增长指标计算函数，计算窗口长度为3的严格增长指标，并将结果追加到原始数据后面
        self.data.ta.increasing(length=3, strict=True, append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名为 "SINC_3"
        self.assertEqual(self.data.columns[-1], "SINC_3")
    # 测试长期趋势计算函数的扩展功能
    def test_long_run_ext(self):
        # 未传递参数，返回原始数据
        self.assertEqual(self.data.ta.long_run(append=True).shape, self.data.shape)
        # 计算快速和慢速指数移动平均线
        fast = self.data.ta.ema(8)
        slow = self.data.ta.ema(21)
        # 计算长期趋势指标，并将结果追加到原始数据后面
        self.data.ta.long_run(fast, slow, append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名为 "LR_2"
        self.assertEqual(self.data.columns[-1], "LR_2")
    # 测试抛物线停止和反转指标计算函数的扩展功能
    def test_psar_ext(self):
        # 计算抛物线停止和反转指标，并将结果追加到原始数据后面
        self.data.ta.psar(append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后四列的列名为指定的抛物线参数
        self.assertEqual(
            list(self.data.columns[-4:]), ["PSARl_0.02_0.2", "PSARs_0.02_0.2", "PSARaf_0.02_0.2", "PSARr_0.02_0.2"])
    # 测试 QStick 指标计算函数的扩展功能
    def test_qstick_ext(self):
        # 计算 QStick 指标，并将结果追加到原始数据后面
        self.data.ta.qstick(append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名为 "QS_10"
        self.assertEqual(self.data.columns[-1], "QS_10")
    # 测试短期趋势计算函数的扩展功能
    def test_short_run_ext(self):
        # 未传递参数，返回原始数据
        self.assertEqual(
            self.data.ta.short_run(append=True).shape, self.data.shape)
        # 计算快速和慢速指数移动平均线
        fast = self.data.ta.ema(8)
        slow = self.data.ta.ema(21)
        # 计算短期趋势指标，并将结果追加到原始数据后面
        self.data.ta.short_run(fast, slow, append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名为 "SR_2"
        self.assertEqual(self.data.columns[-1], "SR_2")
    # 测试 TTM 趋势指标计算函数的扩展功能
    def test_ttm_trend_ext(self):
        # 计算 TTM 趋势指标，并将结果追加到原始数据后面
        self.data.ta.ttm_trend(append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后一列的列名为 "TTM_TRND_6"
        self.assertEqual(list(self.data.columns[-1:]), ["TTM_TRND_6"])
    # 测试涡轮指标计算函数的扩展功能
    def test_vortext_ext(self):
        # 计算涡轮指标，并将结果追加到原始数据后面
        self.data.ta.vortex(append=True)
        # 断言数据类型为 DataFrame
        self.assertIsInstance(self.data, DataFrame)
        # 断言最后两列的列名为 "VTXP_14" 和 "VTXM_14"
        self.assertEqual(list(self.data.columns[-2:]), ["VTXP_14", "VTXM_14"])

# 导入所需模块
from .config import sample_data
from .context import pandas_ta
# 导入 TestCase 类
from unittest import TestCase
# 导入 DataFrame 类
from pandas import DataFrame
# 创建测试类 TestVolatilityExtension，继承自 TestCase 类
class TestVolatilityExtension(TestCase):
    # 在所有测试方法执行之前执行的方法
    @classmethod
    def setUpClass(cls):
        # 设置测试数据
        cls.data = sample_data
    # 在所有测试方法执行之后执行的方法
    @classmethod
    def tearDownClass(cls):
        # 删除测试数据
        del cls.data
    # 在每个测试方法执行之前执行的方法
    def setUp(self): pass
    
    # 在每个测试方法执行之后执行的方法
    def tearDown(self): pass
    # 测试 aberration 方法是否正常工作
    def test_aberration_ext(self):
        # 调用 aberration 方法
        self.data.ta.aberration(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后四列是否等于给定列表
        self.assertEqual(list(self.data.columns[-4:]), ["ABER_ZG_5_15", "ABER_SG_5_15", "ABER_XG_5_15", "ABER_ATR_5_15"])
    # 测试 accbands 方法是否正常工作
    def test_accbands_ext(self):
        # 调用 accbands 方法
        self.data.ta.accbands(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后三列是否等于给定列表
        self.assertEqual(list(self.data.columns[-3:]), ["ACCBL_20", "ACCBM_20", "ACCBU_20"])
    # 测试 atr 方法是否正常工作
    def test_atr_ext(self):
        # 调用 atr 方法
        self.data.ta.atr(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后一列是否等于给定字符串
        self.assertEqual(self.data.columns[-1], "ATRr_14")
    # 测试 bbands 方法是否正常工作
    def test_bbands_ext(self):
        # 调用 bbands 方法
        self.data.ta.bbands(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后五列是否等于给定列表
        self.assertEqual(list(self.data.columns[-5:]), ["BBL_5_2.0", "BBM_5_2.0", "BBU_5_2.0", "BBB_5_2.0", "BBP_5_2.0"])
    # 测试 donchian 方法是否正常工作
    def test_donchian_ext(self):
        # 调用 donchian 方法
        self.data.ta.donchian(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后三列是否等于给定列表
        self.assertEqual(list(self.data.columns[-3:]), ["DCL_20_20", "DCM_20_20", "DCU_20_20"])
    # 测试 kc 方法是否正常工作
    def test_kc_ext(self):
        # 调用 kc 方法
        self.data.ta.kc(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后三列是否等于给定列表
        self.assertEqual(list(self.data.columns[-3:]), ["KCLe_20_2", "KCBe_20_2", "KCUe_20_2"])
    # 测试 massi 方法是否正常工作
    def test_massi_ext(self):
        # 调用 massi 方法
        self.data.ta.massi(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后一列是否等于给定字符串
        self.assertEqual(self.data.columns[-1], "MASSI_9_25")
    # 测试 natr 方法是否正常工作
    def test_natr_ext(self):
        # 调用 natr 方法
        self.data.ta.natr(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后一列是否等于给定字符串
        self.assertEqual(self.data.columns[-1], "NATR_14")
    # 测试 pdist 方法是否正常工作
    def test_pdist_ext(self):
        # 调用 pdist 方法
        self.data.ta.pdist(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后一列是否等于给定字符串
        self.assertEqual(self.data.columns[-1], "PDIST")
    # 测试 rvi 方法是否正常工作
    def test_rvi_ext(self):
        # 调用 rvi 方法
        self.data.ta.rvi(append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后一列是否等于给定字符串
        self.assertEqual(self.data.columns[-1], "RVI_14")
    # 测试 rvi 方法是否正常工作，且使用 refined 参数
    def test_rvi_refined_ext(self):
        # 调用 rvi 方法，传入 refined 参数为 True
        self.data.ta.rvi(refined=True, append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后一列是否等于给定字符串
        self.assertEqual(self.data.columns[-1], "RVIr_14")
    # 测试 rvi 方法是否正常工作，且使用 thirds 参数
    def test_rvi_thirds_ext(self):
        # 调用 rvi 方法，传入 thirds 参数为 True
        self.data.ta.rvi(thirds=True, append=True)
        # 断言 data 是否为 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 断言 data 的最后一列是否等于给定字符串
        self.assertEqual(self.data.columns[-1], "RVIt_14")
    # 测试扩展函数：计算热力指标
    def test_thermo_ext(self):
        # 调用热力指标计算函数，并将结果附加到原始数据上
        self.data.ta.thermo(append=True)
        # 确保结果是 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 检查结果中最后四列的列名是否符合预期
        self.assertEqual(list(self.data.columns[-4:]), ["THERMO_20_2_0.5", "THERMOma_20_2_0.5", "THERMOl_20_2_0.5", "THERMOs_20_2_0.5"])
    # 测试扩展函数：计算真实范围
    def test_true_range_ext(self):
        # 调用真实范围计算函数，并将结果附加到原始数据上
        self.data.ta.true_range(append=True)
        # 确保结果是 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 检查结果中最后一列的列名是否符合预期
        self.assertEqual(self.data.columns[-1], "TRUERANGE_1")
    # 测试扩展函数：计算 UI 指标
    def test_ui_ext(self):
        # 调用 UI 指标计算函数，并将结果附加到原始数据上
        self.data.ta.ui(append=True)
        # 确保结果是 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 检查结果中最后一列的列名是否符合预期
        self.assertEqual(self.data.columns[-1], "UI_14")
        # 调用 UI 指标计算函数，使用 everget 参数，并将结果附加到原始数据上
        self.data.ta.ui(append=True, everget=True)
        # 确保结果是 DataFrame 类型
        self.assertIsInstance(self.data, DataFrame)
        # 检查结果中最后一列的列名是否符合预期
        self.assertEqual(self.data.columns[-1], "UIe_14")