Move main() to its own function

1 files changed, 43 insertions, 33 deletions

diff --git a/addon.py b/addon.py
index 56a7dca..17b316e 100644
--- a/addon.py
+++ b/addon.py
@@ -4,29 +4,28 @@
 from bs4 import BeautifulSoup
 import sys
 import urllib
-import urllib2
+#import urllib2
+import requests
 import urlparse
+import xbmc
 import xbmcgui
 import xbmcplugin
 
+
 KAN_URL = 'http://www.kan.org.il'
 USER_AGENT = 'Kodi'   # I'm optimistic
 
-base_url = sys.argv[0]
-addon_handle = int(sys.argv[1])
-args = urlparse.parse_qs(sys.argv[2][1:])
-mode = args.get('mode', None)
-
 
-def build_url(query):
+def build_url(base_url, query):
     return base_url + '?' + urllib.urlencode(query)
 
 
 def read_url(url):
-    req = urllib2.Request(url)
-    req.add_header('User-Agent', USER_AGENT)
-    response = urllib2.urlopen(req, timeout=100)
-    return response.read()
+    headers = {'user-agent': 'Kodi/plugin.video.kan'}
+    response = requests.get(url, headers=headers)
+    if response.status_code != 200:
+        raise IOError("Invalid URL {}".format(url))
+    return response.content
 
 
 def placeholder_folder(foldername):
@@ -40,37 +39,48 @@ def placeholder_folder(foldername):
 def tvshows_menu():
     tvshows_url = 'http://www.kan.org.il/video/programs.aspx'
     main_page = read_url(tvshows_url)
-    parsed = BeautifulSoup(main_page)
+    parsed = BeautifulSoup(main_page, "lxml")
     anchors = parsed.find_all('a',
                               class_="program_category_link w-inline-block")
     for a in anchors:
-        path = a.href
+        path = a.get('href')
         li = xbmcgui.ListItem(foldername + 'Show' + path)
         xbmcplugin.addDirectoryItem(handle=addon_handle, listitem=li)
     xbmcplugin.endOfDirectory(addon_handle)
 
 
-if mode is None:
-    url = build_url({'mode': 'folder', 'foldername': 'tv-shows'})
-    li = xbmcgui.ListItem(u'תוכניות טלוויזיה')
-    xbmcplugin.addDirectoryItem(handle=addon_handle, url=url,
-                                listitem=li, isFolder=True)
+def main():
+    base_url = sys.argv[0]
+    addon_handle = int(sys.argv[1])
+    args = urlparse.parse_qs(sys.argv[2][1:])
+    mode = args.get('mode', None)
 
-    url = build_url({'mode': 'folder', 'foldername': 'net-shows'})
-    li = xbmcgui.ListItem(u'תוכניות רשת')
-    xbmcplugin.addDirectoryItem(handle=addon_handle, url=url,
-                                listitem=li, isFolder=True)
+    xbmc.log("Tzafrir: base_url: {}, addon_handle: {}, mode: {}".format(base_url, addon_handle, mode), xbmc.LOGINFO)
+    if mode is None:
+        url = build_url(base_url, {'mode': 'folder', 'foldername': 'tv-shows'})
+        li = xbmcgui.ListItem(u'תוכניות טלוויזיה')
+        xbmcplugin.addDirectoryItem(handle=addon_handle, url=url,
+                                    listitem=li, isFolder=True)
 
-    url = build_url({'mode': 'folder', 'foldername': 'new-items'})
-    li = xbmcgui.ListItem(u'קטעים חדשים')
-    xbmcplugin.addDirectoryItem(handle=addon_handle, url=url,
-                                listitem=li, isFolder=True)
+        url = build_url(base_url, {'mode': 'folder', 'foldername': 'net-shows'})
+        li = xbmcgui.ListItem(u'תוכניות רשת')
+        xbmcplugin.addDirectoryItem(handle=addon_handle, url=url,
+                                    listitem=li, isFolder=True)
+
+        url = build_url(base_url, {'mode': 'folder', 'foldername': 'new-items'})
+        li = xbmcgui.ListItem(u'קטעים חדשים')
+        xbmcplugin.addDirectoryItem(handle=addon_handle, url=url,
+                                    listitem=li, isFolder=True)
+
+        xbmcplugin.endOfDirectory(addon_handle)
+
+    elif mode[0] == 'folder':
+        foldername = args['foldername'][0]
+        if foldername == 'tv-shows':
+            tvshows_menu()
+        else:
+            placeholder_folder(foldername)
 
-    xbmcplugin.endOfDirectory(addon_handle)
 
-elif mode[0] == 'folder':
-    foldername = args['foldername'][0]
-    if foldername == 'tv-shows':
-        tvshows_menu()
-    else:
-        placeholder_folder(foldername)
+if __name__ == '__main__':
+    main()