feat: Add live to stream info to match other implementations

pawanpaudel93 · Nov 1, 2023 · 1c74182 · 1c74182
1 parent 305f329
commit 1c74182
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/m3u_parser/m3u_parser.py b/m3u_parser/m3u_parser.py
@@ -199,10 +199,10 @@ async def _get_status(self, stream_link):
                     timeout=self._timeout,
                 ) as response:
                     if response.status == 200:
-                        return "GOOD"
+                        return True
         except:
             pass
-        return "BAD"
+        return False
 
     async def _check_status(self, index):
         stream_info = self._streams_info[index]
@@ -211,7 +211,8 @@ async def _check_status(self, index):
         status_fn = self._parse_config.status_checker.get(scheme)
         if status_fn is None or not callable(status_fn):
             status_fn = self._get_status
-        stream_info["status"] = "GOOD" if await status_fn(stream_url) == "GOOD" else "BAD"
+        stream_info["status"] = "GOOD" if await status_fn(stream_url) == True else "BAD"
+        stream_info["live"] = stream_info["status"] == "GOOD"
         self._streams_info[index] = stream_info
 
     def _check_streams_status(self):
@@ -287,9 +288,10 @@ async def _parse_line(self, line_num: int):
                 status_fn = self._parse_config.status_checker.get(scheme)
                 if status_fn is None or not callable(status_fn):
                     status_fn = self._get_status
-                status = "GOOD" if await status_fn(stream_link) == "GOOD" else "BAD"
+                status = "GOOD" if await status_fn(stream_link) == True else "BAD"
             if self._check_live:
                 info["status"] = status
+                info["live"] = status == "GOOD"
             self._streams_info.append(info)
 
     @staticmethod
@@ -404,6 +406,7 @@ def parse_json(self, data_source: str, config: ParseConfig = ParseConfig()):
                     "country": {"code": stream_info.get("country_code"), "name": stream_info.get("country_name")},
                     "language": {"code": stream_info.get("language_code"), "name": stream_info.get("language_name")},
                     "status": stream_info.get("status") or "BAD",
+                    "live": stream_info.get("status") == "GOOD",
                 }
                 for stream_info in streams_info
                 if type(stream_info) == dict and stream_info.get("url")
@@ -453,6 +456,7 @@ def parse_csv(self, data_source: str, config: ParseConfig = ParseConfig()):
                 "country": {"code": get_value(row, "country_code"), "name": get_value(row, "country_name")},
                 "language": {"code": get_value(row, "language_code"), "name": get_value(row, "language_name")},
                 "status": get_value(row, "status") or "BAD",
+                "live": get_value(row, "status") == "GOOD",
             }
             for row in reader
             if get_value(row, "url")

diff --git a/tests/test_m3uparser.py b/tests/test_m3uparser.py
@@ -66,7 +66,7 @@
 
 
 async def rtsp_checker(url: str):
-    return "GOOD"
+    return True
 
 
 # Fixture to create a temporary M3U file for testing
@@ -195,6 +195,7 @@ def test_save_to_json(self, temp_m3u_file, tmpdir):
         parser.to_file(str(json_file), format="json")
         print(json_file)
         assert os.path.exists(str(json_file))
+        os.remove(str(json_file))
 
     # Test saving to CSV file
     def test_save_to_csv(self, temp_m3u_file, tmpdir):
@@ -204,6 +205,7 @@ def test_save_to_csv(self, temp_m3u_file, tmpdir):
         parser.to_file(str(csv_file), format="csv")
         print(csv_file)
         assert os.path.exists(str(csv_file))
+        os.remove(str(csv_file))
 
     # Test filtering by category
     def test_filter_by_category(self, temp_m3u_file):