How to use b2luigi - 10 common examples

import b2luigi
import random


class MyNumberTask(b2luigi.Task):
    some_parameter = b2luigi.IntParameter()

    def output(self):
        return b2luigi.LocalTarget(f"results/output_file_{self.some_parameter}.txt")

    def run(self):
        random_number = random.random()
        with self.output().open("w") as f:
            f.write(f"{random_number}\n")


if __name__ == "__main__":
    b2luigi.set_setting("result_dir", "results")
    b2luigi.process([MyNumberTask(some_parameter=i) for i in range(100)], 
                    workers=200)

def test_requires(self):
        class TaskA(b2luigi.Task):
            some_parameter = b2luigi.IntParameter()
            some_other_parameter = b2luigi.IntParameter()

            def output(self):
                yield self.add_to_output("test.txt")

        @b2luigi.requires(TaskA, some_parameter=3)
        class TaskB(b2luigi.Task):
            another_parameter = b2luigi.IntParameter()

            def output(self):
                yield self.add_to_output("out.dat")

        task = TaskB(some_other_parameter=1, another_parameter=42)
        self.assertEqual(sorted(task.get_param_names()), ["another_parameter", "some_other_parameter"])
        self.assertEqual(task.another_parameter, 42)
        self.assertEqual(task.some_other_parameter, 1)

def test_dependencies(self):
        class TaskA(b2luigi.Task):
            some_parameter = b2luigi.IntParameter()

            def output(self):
                yield self.add_to_output("file_a")

        @b2luigi.requires(TaskA)
        class TaskB(b2luigi.Task):
            def output(self):
                yield self.add_to_output("file_b")

        task = TaskB(some_parameter=42)

        self.assertEqual(get_filled_params(task), {"some_parameter": 42})
        self.assertEqual(len(task._get_input_targets("file_a")), 1)
        self.assertEqual(len(task.get_input_file_names("file_a")), 1)
        self.assertEqual(len(task.get_input_file_names().keys()), 1)
        self.assertEqual(task._get_input_targets("file_a")[0].path, task.get_input_file_names("file_a")[0])

def test_requires(self):
        class TaskA(b2luigi.Task):
            some_parameter = b2luigi.IntParameter()
            some_other_parameter = b2luigi.IntParameter()

            def output(self):
                yield self.add_to_output("test.txt")

        @b2luigi.requires(TaskA, some_parameter=3)
        class TaskB(b2luigi.Task):
            another_parameter = b2luigi.IntParameter()

            def output(self):
                yield self.add_to_output("out.dat")

        task = TaskB(some_other_parameter=1, another_parameter=42)
        self.assertEqual(sorted(task.get_param_names()), ["another_parameter", "some_other_parameter"])
        self.assertEqual(task.another_parameter, 42)
        self.assertEqual(task.some_other_parameter, 1)

def test_file_path_usage(self):
        class TaskA(b2luigi.Task):
            some_parameter = b2luigi.IntParameter()

            def output(self):
                yield self.add_to_output("file_a")
                yield self.add_to_output("file_b")
        
        task = TaskA(some_parameter=3)

        b2luigi.set_setting("result_dir", "results/some_crazy_path")
        
        self.assertEqual(get_filled_params(task), {"some_parameter": 3})
        self.assertFalse(task.get_input_file_names())
        self.assertRaises(KeyError, lambda: task._get_input_targets("some_file"))
        self.assertEqual(task._get_output_target("file_a").path, task.get_output_file_name("file_a"))
        self.assertIn("file_a", task.get_output_file_name("file_a"))
        self.assertIn("file_b", task.get_output_file_name("file_b"))
        self.assertIn("some_parameter=3", task.get_output_file_name("file_a"))

def test_set_by_file(self):
        with open("settings.json", "w") as f:
            json.dump({"my_setting": "my file value"}, f)

        self.assertEqual("my file value", b2luigi.get_setting("my_setting"))

        b2luigi.set_setting("my_setting", "my value")

        self.assertEqual("my value", b2luigi.get_setting("my_setting"))

def run(self):
        print("Hello!")
        with open(self.get_output_file_name("some_file.txt"), "w") as f:
            f.write("Done")

        print("Bye!")
        import sys
        sys.stdout.flush()
        os.kill(os.getpid(), 11)

        with open(self.get_output_file_name("some_other_file.txt"), "w") as f:
            f.write("Done")


if __name__ == "__main__":
    b2luigi.set_setting("result_dir", "results")
    b2luigi.process(MyTask())

def test_deprecated_settings(self):
        self.assertRaises(ValueError, b2luigi.get_setting, key="my_setting", 
                          deprecated_keys=["my_old_setting"])

        b2luigi.set_setting("my_old_setting", "my value")

        with warnings.catch_warnings(record=True) as w:
            self.assertEqual("my value", b2luigi.get_setting("my_setting", 
                            deprecated_keys=["my_old_setting"]))

            self.assertEqual(len(w), 1)
            self.assertIsInstance(w[-1].message, DeprecatedSettingsWarning)
            self.assertIn("deprecated", str(w[-1].message))

        b2luigi.set_setting("my_setting", "my new_value")

        with warnings.catch_warnings(record=True) as w:
            self.assertEqual("my new_value", b2luigi.get_setting("my_setting", default="default", 
                            deprecated_keys=["my_old_setting"]))

            self.assertEqual(len(w), 0)

def process(self):
        print("Hello!")
        with open(self.get_output_file_name("some_file.txt"), "w") as f:
            f.write("Done")

        print("Bye!")
        import sys
        sys.stdout.flush()
        os.kill(os.getpid(), 11)

        with open(self.get_output_file_name("some_other_file.txt"), "w") as f:
            f.write("Done")


if __name__ == "__main__":
    b2luigi.set_setting("result_dir", "results")
    b2luigi.process(MyTask())

def test_file_path_usage(self):
        class TaskA(b2luigi.Task):
            some_parameter = b2luigi.IntParameter()

            def output(self):
                yield self.add_to_output("file_a")
                yield self.add_to_output("file_b")
        
        task = TaskA(some_parameter=3)

        b2luigi.set_setting("result_dir", "results/some_crazy_path")
        
        self.assertEqual(get_filled_params(task), {"some_parameter": 3})
        self.assertFalse(task.get_input_file_names())
        self.assertRaises(KeyError, lambda: task._get_input_targets("some_file"))
        self.assertEqual(task._get_output_target("file_a").path, task.get_output_file_name("file_a"))
        self.assertIn("file_a", task.get_output_file_name("file_a"))
        self.assertIn("file_b", task.get_output_file_name("file_b"))
        self.assertIn("some_parameter=3", task.get_output_file_name("file_a"))
        self.assertIn("some_crazy_path", task.get_output_file_name("file_a"))